screen: update tresholds on text
This commit is contained in:
52
screen.go
52
screen.go
@@ -67,21 +67,17 @@ func (d *Device) GetScreenResolution() (int, int, error) {
|
|||||||
return d.Screenx, d.Screeny, nil
|
return d.Screenx, d.Screeny, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// Get the text on screen in a certain area using gotesseract
|
func (d *Device) GetScreenText(xStart int, xEnd int, yStart int, yEnd int, whitelist string, multiline bool, textPolarity string) (string, error) {
|
||||||
func (d *Device) GetScreenText(xStart int, xEnd int, yStart int, yEnd int, whitelist string, multiline bool) (string, error) {
|
|
||||||
// Take screenshot in memory and get bytes
|
|
||||||
data, err := d.RunCommand("exec-out screencap -p")
|
data, err := d.RunCommand("exec-out screencap -p")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", fmt.Errorf("failed to take screenshot: %w", err)
|
return "", fmt.Errorf("failed to take screenshot: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Decode image from bytes
|
|
||||||
img, err := png.Decode(bytes.NewReader(data))
|
img, err := png.Decode(bytes.NewReader(data))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", fmt.Errorf("failed to decode PNG: %w", err)
|
return "", fmt.Errorf("failed to decode PNG: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Crop region
|
|
||||||
rect := image.Rect(0, 0, xEnd-xStart, yEnd-yStart)
|
rect := image.Rect(0, 0, xEnd-xStart, yEnd-yStart)
|
||||||
cropped := image.NewRGBA(rect)
|
cropped := image.NewRGBA(rect)
|
||||||
for y := yStart; y < yEnd; y++ {
|
for y := yStart; y < yEnd; y++ {
|
||||||
@@ -90,27 +86,57 @@ func (d *Device) GetScreenText(xStart int, xEnd int, yStart int, yEnd int, white
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Convert to grayscale and apply threshold
|
// Convert to grayscale and compute average brightness
|
||||||
threshold := uint8(160)
|
gray := image.NewGray(rect)
|
||||||
|
var totalBrightness int
|
||||||
|
for y := 0; y < rect.Dy(); y++ {
|
||||||
|
for x := 0; x < rect.Dx(); x++ {
|
||||||
|
grayVal := color.GrayModel.Convert(cropped.At(x, y)).(color.Gray)
|
||||||
|
gray.Set(x, y, grayVal)
|
||||||
|
totalBrightness += int(grayVal.Y)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
avgBrightness := totalBrightness / (rect.Dx() * rect.Dy())
|
||||||
|
|
||||||
|
// Decide polarity
|
||||||
|
useLightText := false
|
||||||
|
switch textPolarity {
|
||||||
|
case "light":
|
||||||
|
useLightText = true
|
||||||
|
case "dark":
|
||||||
|
useLightText = false
|
||||||
|
case "auto":
|
||||||
|
useLightText = avgBrightness < 128
|
||||||
|
default:
|
||||||
|
return "", fmt.Errorf("invalid textPolarity: must be 'dark', 'light', or 'auto'")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Binarize based on polarity
|
||||||
binarized := image.NewGray(rect)
|
binarized := image.NewGray(rect)
|
||||||
for y := 0; y < rect.Dy(); y++ {
|
for y := 0; y < rect.Dy(); y++ {
|
||||||
for x := 0; x < rect.Dx(); x++ {
|
for x := 0; x < rect.Dx(); x++ {
|
||||||
gray := color.GrayModel.Convert(cropped.At(x, y)).(color.Gray)
|
grayVal := gray.GrayAt(x, y).Y
|
||||||
if gray.Y > threshold {
|
if useLightText {
|
||||||
binarized.SetGray(x, y, color.Gray{Y: 255})
|
if int(grayVal) > avgBrightness+10 {
|
||||||
|
binarized.SetGray(x, y, color.Gray{Y: 0})
|
||||||
|
} else {
|
||||||
|
binarized.SetGray(x, y, color.Gray{Y: 255})
|
||||||
|
}
|
||||||
} else {
|
} else {
|
||||||
binarized.SetGray(x, y, color.Gray{Y: 0})
|
if int(grayVal) < avgBrightness-10 {
|
||||||
|
binarized.SetGray(x, y, color.Gray{Y: 0})
|
||||||
|
} else {
|
||||||
|
binarized.SetGray(x, y, color.Gray{Y: 255})
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Encode binarized image to PNG
|
|
||||||
var buf bytes.Buffer
|
var buf bytes.Buffer
|
||||||
if err := png.Encode(&buf, binarized); err != nil {
|
if err := png.Encode(&buf, binarized); err != nil {
|
||||||
return "", fmt.Errorf("failed to encode image: %w", err)
|
return "", fmt.Errorf("failed to encode image: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Call tesseract
|
|
||||||
client := gosseract.NewClient()
|
client := gosseract.NewClient()
|
||||||
defer client.Close()
|
defer client.Close()
|
||||||
client.SetImageFromBytes(buf.Bytes())
|
client.SetImageFromBytes(buf.Bytes())
|
||||||
|
|||||||
Reference in New Issue
Block a user