httprunner/hrp/pkg/uixt/ocr_on.go

//go:build ocr

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
}

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).Msgf("FindText success")
	x, y, width, height = dExt.MappingToRectInUIKit(rect)
	return
}