refactor: vedem ocr get texts

hrp/pkg/uixt/ocr_vedem.go

@@ -0,0 +1,246 @@
+package uixt
+
+import (
+	"bytes"
+	"fmt"
+	"image"
+	"io/ioutil"
+	"mime/multipart"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+
+	"github.com/rs/zerolog/log"
+
+	"github.com/httprunner/httprunner/v4/hrp/internal/env"
+	"github.com/httprunner/httprunner/v4/hrp/internal/json"
+)
+
+var client = &http.Client{
+	Timeout: time.Second * 10,
+}
+
+type OCRResult struct {
+	Text   string   `json:"text"`
+	Points []PointF `json:"points"`
+}
+
+type ResponseOCR struct {
+	Code      int         `json:"code"`
+	Message   string      `json:"message"`
+	OCRResult []OCRResult `json:"ocrResult"`
+}
+
+type veDEMOCRService struct{}
+
+func (s *veDEMOCRService) getOCRResult(imageBuf []byte) ([]OCRResult, error) {
+	bodyBuf := &bytes.Buffer{}
+	bodyWriter := multipart.NewWriter(bodyBuf)
+	bodyWriter.WriteField("withDet", "true")
+	// bodyWriter.WriteField("timestampOnly", "true")
+
+	formWriter, err := bodyWriter.CreateFormFile("image", "screenshot.png")
+	if err != nil {
+		return nil, fmt.Errorf("create form file error: %v", err)
+	}
+	_, err = formWriter.Write(imageBuf)
+	if err != nil {
+		return nil, fmt.Errorf("write form error: %v", err)
+	}
+
+	err = bodyWriter.Close()
+	if err != nil {
+		return nil, fmt.Errorf("close body writer error: %v", err)
+	}
+
+	if env.VEDEM_OCR_URL == "" {
+		log.Error().Msg("VEDEM_OCR_URL env missed for OCR service")
+		os.Exit(1)
+	}
+
+	req, err := http.NewRequest("POST", env.VEDEM_OCR_URL, bodyBuf)
+	if err != nil {
+		return nil, fmt.Errorf("construct request error: %v", err)
+	}
+
+	req.Header.Add("Content-Type", bodyWriter.FormDataContentType())
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("http reqeust OCR server error: %v", err)
+	}
+	defer resp.Body.Close()
+
+	results, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("read response body error: %v", err)
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("unexpected response status code: %d, results: %v", resp.StatusCode, string(results))
+	}
+
+	var ocrResult ResponseOCR
+	err = json.Unmarshal(results, &ocrResult)
+	if err != nil {
+		return nil, fmt.Errorf("json unmarshal response body error: %v", err)
+	}
+
+	return ocrResult.OCRResult, nil
+}
+
+func (s *veDEMOCRService) FindText(text string, imageBuf []byte, index ...int) (rect image.Rectangle, err error) {
+	if len(index) == 0 {
+		index = []int{0} // index not specified
+	}
+
+	ocrResults, err := s.getOCRResult(imageBuf)
+	if err != nil {
+		log.Error().Err(err).Msg("getOCRResult failed")
+		return
+	}
+
+	var rects []image.Rectangle
+	var ocrTexts []string
+	for _, ocrResult := range ocrResults {
+		ocrTexts = append(ocrTexts, ocrResult.Text)
+
+		// not contains text
+		if !strings.Contains(ocrResult.Text, text) {
+			continue
+		}
+
+		rect = image.Rectangle{
+			// ocrResult.Points 顺序：左上 -> 右上 -> 右下 -> 左下
+			Min: image.Point{
+				X: int(ocrResult.Points[0].X),
+				Y: int(ocrResult.Points[0].Y),
+			},
+			Max: image.Point{
+				X: int(ocrResult.Points[2].X),
+				Y: int(ocrResult.Points[2].Y),
+			},
+		}
+		rects = append(rects, rect)
+
+		// contains text while not match exactly
+		if ocrResult.Text != text {
+			continue
+		}
+
+		// match exactly, and not specify index, return the first one
+		if index[0] == 0 {
+			return rect, nil
+		}
+	}
+
+	if len(rects) == 0 {
+		return image.Rectangle{},
+			fmt.Errorf("text %s not found in %v", text, ocrTexts)
+	}
+
+	// get index
+	idx := index[0]
+	if idx > 0 {
+		// NOTICE: index start from 1
+		idx = idx - 1
+	} else if idx < 0 {
+		idx = len(rects) + idx
+	}
+
+	// index out of range
+	if idx >= len(rects) {
+		return image.Rectangle{}, fmt.Errorf("text %s found %d, index %d out of range",
+			text, len(rects), idx)
+	}
+
+	return rects[idx], nil
+}
+
+func (s *veDEMOCRService) FindTexts(texts []string, imageBuf []byte) (rects []image.Rectangle, err error) {
+	ocrResults, err := s.getOCRResult(imageBuf)
+	if err != nil {
+		log.Error().Err(err).Msg("getOCRResult failed")
+		return
+	}
+
+	for _, text := range texts {
+		var found bool
+		for _, ocrResult := range ocrResults {
+			// not contains text
+			if !strings.Contains(ocrResult.Text, text) {
+				continue
+			}
+
+			found = true
+			rect := image.Rectangle{
+				// ocrResult.Points 顺序：左上 -> 右上 -> 右下 -> 左下
+				Min: image.Point{
+					X: int(ocrResult.Points[0].X),
+					Y: int(ocrResult.Points[0].Y),
+				},
+				Max: image.Point{
+					X: int(ocrResult.Points[2].X),
+					Y: int(ocrResult.Points[2].Y),
+				},
+			}
+			rects = append(rects, rect)
+			break
+		}
+		if !found {
+			rects = append(rects, image.Rectangle{})
+		}
+	}
+
+	return rects, nil
+}
+
+type OCRService interface {
+	FindText(text string, imageBuf []byte, index ...int) (rect image.Rectangle, err error)
+}
+
+func (dExt *DriverExt) FindTextByOCR(ocrText string, index ...int) (x, y, width, height float64, err error) {
+	var bufSource *bytes.Buffer
+	if bufSource, err = dExt.takeScreenShot(); err != nil {
+		err = fmt.Errorf("takeScreenShot error: %v", err)
+		return
+	}
+
+	service := &veDEMOCRService{}
+	rect, err := service.FindText(ocrText, bufSource.Bytes(), index...)
+	if err != nil {
+		log.Warn().Msgf("FindText failed: %s", err.Error())
+		err = fmt.Errorf("FindText failed: %v", err)
+		return
+	}
+
+	log.Info().Str("ocrText", ocrText).
+		Interface("rect", rect).Msgf("FindTextByOCR success")
+	x, y, width, height = dExt.MappingToRectInUIKit(rect)
+	return
+}
+
+func (dExt *DriverExt) FindTextsByOCR(ocrTexts []string) (points [][]float64, err error) {
+	var bufSource *bytes.Buffer
+	if bufSource, err = dExt.takeScreenShot(); err != nil {
+		err = fmt.Errorf("takeScreenShot error: %v", err)
+		return
+	}
+
+	service := &veDEMOCRService{}
+	rects, err := service.FindTexts(ocrTexts, bufSource.Bytes())
+	if err != nil {
+		log.Warn().Msgf("FindTexts failed: %s", err.Error())
+		err = fmt.Errorf("FindTexts failed: %v", err)
+		return
+	}
+
+	log.Info().Interface("ocrTexts", ocrTexts).
+		Interface("rects", rects).Msgf("FindTextsByOCR success")
+	for _, rect := range rects {
+		x, y, width, height := dExt.MappingToRectInUIKit(rect)
+		points = append(points, []float64{x, y, width, height})
+	}
+
+	return
+}