168 lines
4.7 KiB
Go
168 lines
4.7 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"image"
|
|
"image/color"
|
|
|
|
"gocv.io/x/gocv"
|
|
)
|
|
|
|
// ContourInfo tracks a detected contour and whether it was used
|
|
type ContourInfo struct {
|
|
rect image.Rectangle
|
|
used bool
|
|
}
|
|
|
|
// detectScreenWidth finds the screen width via contours (returns 0 on failure)
|
|
// If debugPath is not empty, saves an annotated image showing the detection
|
|
func detectScreenWidth(frame gocv.Mat, debugPath string) int {
|
|
gray := gocv.NewMat()
|
|
defer gray.Close()
|
|
gocv.CvtColor(frame, &gray, gocv.ColorBGRToGray)
|
|
|
|
thresh := gocv.NewMat()
|
|
defer thresh.Close()
|
|
gocv.Threshold(gray, &thresh, 170, 255, gocv.ThresholdBinary)
|
|
|
|
contours := gocv.FindContours(thresh, gocv.RetrievalExternal, gocv.ChainApproxSimple)
|
|
defer contours.Close()
|
|
|
|
if contours.Size() == 0 {
|
|
return 0
|
|
}
|
|
|
|
// Collect all contours and track which are used
|
|
var allContours []ContourInfo
|
|
|
|
screenMinX := 999999
|
|
screenMaxX := 0
|
|
for i := 0; i < contours.Size(); i++ {
|
|
rect := gocv.BoundingRect(contours.At(i))
|
|
used := rect.Dx() > 50 || rect.Dy() > 50
|
|
allContours = append(allContours, ContourInfo{rect: rect, used: used})
|
|
|
|
if used {
|
|
if rect.Min.X < screenMinX {
|
|
screenMinX = rect.Min.X
|
|
}
|
|
if rect.Max.X > screenMaxX {
|
|
screenMaxX = rect.Max.X
|
|
}
|
|
}
|
|
}
|
|
|
|
screenWidth := screenMaxX - screenMinX
|
|
if screenWidth < 100 {
|
|
return 0
|
|
}
|
|
|
|
// Save debug image if requested
|
|
if debugPath != "" {
|
|
saveScreenWidthDebug(frame, allContours, screenMinX, screenMaxX, screenWidth, debugPath)
|
|
}
|
|
|
|
return screenWidth
|
|
}
|
|
|
|
// saveScreenWidthDebug creates an annotated image showing screen width detection
|
|
func saveScreenWidthDebug(frame gocv.Mat, contours []ContourInfo, minX, maxX, width int, path string) {
|
|
debug := frame.Clone()
|
|
defer debug.Close()
|
|
|
|
// Draw all contours
|
|
for i, c := range contours {
|
|
if c.used {
|
|
// Used contours in green
|
|
green := color.RGBA{0, 255, 0, 255}
|
|
gocv.Rectangle(&debug, c.rect, green, 2)
|
|
gocv.PutText(&debug, fmt.Sprintf("#%d", i),
|
|
image.Pt(c.rect.Min.X, c.rect.Min.Y-5),
|
|
gocv.FontHersheyPlain, 1.0, green, 1)
|
|
} else {
|
|
// Unused contours in gray
|
|
gray := color.RGBA{100, 100, 100, 255}
|
|
gocv.Rectangle(&debug, c.rect, gray, 1)
|
|
}
|
|
}
|
|
|
|
// Draw vertical lines at min and max X
|
|
red := color.RGBA{255, 0, 0, 255}
|
|
gocv.Line(&debug, image.Pt(minX, 0), image.Pt(minX, debug.Rows()), red, 3)
|
|
gocv.Line(&debug, image.Pt(maxX, 0), image.Pt(maxX, debug.Rows()), red, 3)
|
|
|
|
// Add text annotations
|
|
yellow := color.RGBA{255, 255, 0, 255}
|
|
gocv.PutText(&debug, fmt.Sprintf("Screen MinX: %d", minX),
|
|
image.Pt(minX+10, 30), gocv.FontHersheyDuplex, 0.8, yellow, 2)
|
|
gocv.PutText(&debug, fmt.Sprintf("Screen MaxX: %d", maxX),
|
|
image.Pt(maxX-200, 30), gocv.FontHersheyDuplex, 0.8, yellow, 2)
|
|
gocv.PutText(&debug, fmt.Sprintf("Screen Width: %dpx", width),
|
|
image.Pt(minX+10, 60), gocv.FontHersheyDuplex, 0.8, yellow, 2)
|
|
gocv.PutText(&debug, fmt.Sprintf("Total contours: %d", len(contours)),
|
|
image.Pt(minX+10, 90), gocv.FontHersheyDuplex, 0.8, yellow, 2)
|
|
|
|
gocv.IMWrite(path, debug)
|
|
fmt.Printf(" Saved width detection debug: %s\n", path)
|
|
}
|
|
|
|
// applyScale resizes frame by the given scale factor
|
|
func applyScale(frame gocv.Mat, scale float64) gocv.Mat {
|
|
if scale == 1.0 {
|
|
return frame // No scaling needed
|
|
}
|
|
|
|
newWidth := int(float64(frame.Cols()) * scale)
|
|
newHeight := int(float64(frame.Rows()) * scale)
|
|
|
|
scaled := gocv.NewMat()
|
|
gocv.Resize(frame, &scaled, image.Pt(newWidth, newHeight), 0, 0, gocv.InterpolationLinear)
|
|
|
|
return scaled
|
|
}
|
|
|
|
// normalizeToWidth resizes frame to make it targetWidth pixels wide
|
|
func normalizeToWidth(frame gocv.Mat, targetWidth int) gocv.Mat {
|
|
if frame.Cols() == targetWidth {
|
|
return frame // Already correct size
|
|
}
|
|
|
|
scale := float64(targetWidth) / float64(frame.Cols())
|
|
newWidth := targetWidth
|
|
newHeight := int(float64(frame.Rows()) * scale)
|
|
|
|
normalized := gocv.NewMat()
|
|
gocv.Resize(frame, &normalized, image.Pt(newWidth, newHeight), 0, 0, gocv.InterpolationLinear)
|
|
|
|
return normalized
|
|
}
|
|
|
|
// detectLayoutWithNormalization performs full layout detection with screen width normalization
|
|
// Returns: layout, scale factor, error
|
|
func detectLayoutWithNormalization(frame gocv.Mat) (*ScreenLayout, float64, error) {
|
|
// 1. Detect screen width
|
|
screenWidth := detectScreenWidth(frame, "")
|
|
if screenWidth == 0 {
|
|
return nil, 0, fmt.Errorf("failed to detect screen width")
|
|
}
|
|
|
|
// 2. Calculate and apply scale
|
|
const targetWidth = 860
|
|
scale := float64(targetWidth) / float64(screenWidth)
|
|
normalized := applyScale(frame, scale)
|
|
if normalized.Ptr() != frame.Ptr() {
|
|
defer normalized.Close()
|
|
}
|
|
|
|
// 3. Detect layout on normalized frame
|
|
layout, rescaled, err := detectScreenLayoutAreas(normalized)
|
|
if !rescaled.Empty() {
|
|
rescaled.Close()
|
|
}
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
|
|
return layout, scale, nil
|
|
}
|