531 lines
15 KiB
Go
531 lines
15 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"image"
|
|
"image/color"
|
|
"math"
|
|
"sort"
|
|
|
|
"gocv.io/x/gocv"
|
|
)
|
|
|
|
// Global DEBUG flag is declared in test_rotation.go
|
|
|
|
// Band represents a vertical band of contours
|
|
type Band struct {
|
|
minX int
|
|
maxX int
|
|
minY int
|
|
maxY int
|
|
}
|
|
|
|
// DetectBands finds and groups contours into vertical bands
|
|
// Only returns bands with height >= minHeight pixels
|
|
func DetectBands(binary gocv.Mat, minHeight int) []Band {
|
|
type Box struct {
|
|
minX int
|
|
maxX int
|
|
minY int
|
|
maxY int
|
|
}
|
|
|
|
// Find contours
|
|
contours := gocv.FindContours(binary, gocv.RetrievalExternal, gocv.ChainApproxSimple)
|
|
|
|
// Get all bounding boxes
|
|
boxes := []Box{}
|
|
for i := 0; i < contours.Size(); i++ {
|
|
rect := gocv.BoundingRect(contours.At(i))
|
|
boxes = append(boxes, Box{
|
|
minX: rect.Min.X,
|
|
maxX: rect.Max.X,
|
|
minY: rect.Min.Y,
|
|
maxY: rect.Max.Y,
|
|
})
|
|
}
|
|
|
|
// Sort boxes by minY
|
|
sort.Slice(boxes, func(i, j int) bool {
|
|
return boxes[i].minY < boxes[j].minY
|
|
})
|
|
|
|
// Group overlapping/adjacent boxes into vertical bands
|
|
bands := []Band{}
|
|
|
|
for _, box := range boxes {
|
|
// Skip boxes shorter than minHeight
|
|
height := box.maxY - box.minY
|
|
if height < minHeight {
|
|
continue
|
|
}
|
|
|
|
if len(bands) == 0 {
|
|
// First band
|
|
bands = append(bands, Band{
|
|
minX: box.minX,
|
|
maxX: box.maxX,
|
|
minY: box.minY,
|
|
maxY: box.maxY,
|
|
})
|
|
} else {
|
|
lastBand := &bands[len(bands)-1]
|
|
|
|
// Check if this box overlaps or is close to last band (within 5px)
|
|
if box.minY <= lastBand.maxY+5 {
|
|
// Merge: extend the band in both X and Y
|
|
if box.minX < lastBand.minX {
|
|
lastBand.minX = box.minX
|
|
}
|
|
if box.maxX > lastBand.maxX {
|
|
lastBand.maxX = box.maxX
|
|
}
|
|
if box.maxY > lastBand.maxY {
|
|
lastBand.maxY = box.maxY
|
|
}
|
|
} else {
|
|
// New band
|
|
bands = append(bands, Band{
|
|
minX: box.minX,
|
|
maxX: box.maxX,
|
|
minY: box.minY,
|
|
maxY: box.maxY,
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
return bands
|
|
}
|
|
|
|
// findBaseline analyzes the baseline in the graph band
|
|
// Returns left point, right point, width, number of sections, and error
|
|
func findBaseline(binary gocv.Mat, graphBand Band) (*image.Point, *image.Point, int, int, error) {
|
|
// Extract just the graph band region
|
|
graphRegion := binary.Region(image.Rect(0, graphBand.minY, binary.Cols(), graphBand.maxY+1))
|
|
|
|
// Find all contours in graph region
|
|
contours := gocv.FindContours(graphRegion, gocv.RetrievalExternal, gocv.ChainApproxSimple)
|
|
|
|
// Get sections (contours) with width >= 5px
|
|
type Section struct {
|
|
minX int
|
|
maxX int
|
|
minY int
|
|
maxY int
|
|
}
|
|
sections := []Section{}
|
|
|
|
for i := 0; i < contours.Size(); i++ {
|
|
rect := gocv.BoundingRect(contours.At(i))
|
|
width := rect.Max.X - rect.Min.X
|
|
height := rect.Max.Y - rect.Min.Y
|
|
// Filter: width >= 5px AND height >= 80px
|
|
if width >= 5 && height >= 80 {
|
|
sections = append(sections, Section{
|
|
minX: rect.Min.X,
|
|
maxX: rect.Max.X,
|
|
minY: rect.Min.Y + graphBand.minY, // Adjust to full image coordinates
|
|
maxY: rect.Max.Y + graphBand.minY,
|
|
})
|
|
}
|
|
}
|
|
|
|
// Check if we have exactly 2 sections
|
|
if len(sections) != 2 {
|
|
return nil, nil, 0, len(sections), fmt.Errorf("found %d sections (need exactly 2)", len(sections))
|
|
}
|
|
|
|
// Sort sections by X position (left to right)
|
|
sort.Slice(sections, func(i, j int) bool {
|
|
return sections[i].minX < sections[j].minX
|
|
})
|
|
|
|
// Horizontal width (left edge of left section to right edge of right section)
|
|
horizontalWidth := sections[1].maxX - sections[0].minX
|
|
|
|
// Find rotation points
|
|
// Left section: lowest white pixel 3px from left edge
|
|
leftSection := sections[0]
|
|
leftX := leftSection.minX + 3
|
|
leftPoint := findLowestWhitePixel(binary, leftX, graphBand.minY, graphBand.maxY)
|
|
|
|
// Right section: lowest white pixel 3px from right edge
|
|
rightSection := sections[1]
|
|
rightX := rightSection.maxX - 3
|
|
rightPoint := findLowestWhitePixel(binary, rightX, graphBand.minY, graphBand.maxY)
|
|
|
|
if leftPoint == nil || rightPoint == nil {
|
|
return nil, nil, 0, len(sections), fmt.Errorf("could not find rotation endpoints")
|
|
}
|
|
|
|
return leftPoint, rightPoint, horizontalWidth, len(sections), nil
|
|
}
|
|
|
|
// findLowestWhitePixel finds the lowest (max Y) white pixel at given X within Y range
|
|
func findLowestWhitePixel(binary gocv.Mat, x int, minY int, maxY int) *image.Point {
|
|
lowestY := -1
|
|
for y := minY; y <= maxY; y++ {
|
|
if binary.GetUCharAt(y, x) == 255 {
|
|
lowestY = y
|
|
}
|
|
}
|
|
if lowestY == -1 {
|
|
return nil
|
|
}
|
|
return &image.Point{X: x, Y: lowestY}
|
|
}
|
|
|
|
// DetectionResult contains rotation and width detection results
|
|
type DetectionResult struct {
|
|
Rotation float64
|
|
Width int
|
|
Success bool
|
|
}
|
|
|
|
// PreprocessFrame takes a raw frame and prepares it for detection
|
|
// Returns the processed binary image
|
|
func PreprocessFrame(frame gocv.Mat) gocv.Mat {
|
|
// Crop top 68 pixels (timestamp)
|
|
cropped := frame.Region(image.Rect(0, 68, frame.Cols(), frame.Rows()))
|
|
|
|
// Rotate 90° clockwise
|
|
rotated := gocv.NewMat()
|
|
gocv.Rotate(cropped, &rotated, gocv.Rotate90Clockwise)
|
|
|
|
// Convert to grayscale
|
|
gray := gocv.NewMat()
|
|
gocv.CvtColor(rotated, &gray, gocv.ColorBGRToGray)
|
|
rotated.Close()
|
|
|
|
// Threshold to binary
|
|
binary := gocv.NewMat()
|
|
gocv.Threshold(gray, &binary, 170, 255, gocv.ThresholdBinary)
|
|
gray.Close()
|
|
|
|
return binary
|
|
}
|
|
|
|
// DetectRotationAndWidth analyzes a binary frame and returns rotation angle and width
|
|
func DetectRotationAndWidth(binary gocv.Mat) DetectionResult {
|
|
// Detect bands (using 80px minimum height for original scale)
|
|
bands := DetectBands(binary, 80)
|
|
|
|
// Check if we have enough bands
|
|
if len(bands) < 2 {
|
|
if DEBUG {
|
|
fmt.Printf("DEBUG: Detection failed - only %d bands found (need at least 2)\n", len(bands))
|
|
}
|
|
return DetectionResult{Success: false}
|
|
}
|
|
|
|
// Find TALLEST band (graph/waveform display)
|
|
tallestIdx := 0
|
|
tallestHeight := 0
|
|
for i, band := range bands {
|
|
height := band.maxY - band.minY
|
|
if height > tallestHeight {
|
|
tallestHeight = height
|
|
tallestIdx = i
|
|
}
|
|
}
|
|
|
|
// Check if there's a band after the tallest (for digits)
|
|
if tallestIdx >= len(bands)-1 {
|
|
if DEBUG {
|
|
fmt.Printf("DEBUG: Detection failed - tallest band is last (idx=%d, total=%d)\n", tallestIdx, len(bands))
|
|
}
|
|
return DetectionResult{Success: false}
|
|
}
|
|
|
|
graphBand := bands[tallestIdx] // Tallest = graph/waveform
|
|
|
|
// Find baseline
|
|
leftPoint, rightPoint, width, sections, err := findBaseline(binary, graphBand)
|
|
if err != nil || sections != 2 {
|
|
if DEBUG {
|
|
if err != nil {
|
|
fmt.Printf("DEBUG: Detection failed - baseline error: %v\n", err)
|
|
} else {
|
|
fmt.Printf("DEBUG: Detection failed - found %d sections (need exactly 2)\n", sections)
|
|
}
|
|
}
|
|
return DetectionResult{Success: false}
|
|
}
|
|
|
|
// Calculate rotation angle
|
|
deltaY := float64(rightPoint.Y - leftPoint.Y)
|
|
deltaX := float64(rightPoint.X - leftPoint.X)
|
|
rotation := math.Atan2(deltaY, deltaX) * 180 / math.Pi
|
|
|
|
return DetectionResult{
|
|
Rotation: rotation,
|
|
Width: width,
|
|
Success: true,
|
|
}
|
|
}
|
|
|
|
// VisualizeBands draws bands on an image with labels (only in DEBUG mode)
|
|
func VisualizeBands(binary gocv.Mat, bands []Band, tallestIdx int, filename string) {
|
|
if !DEBUG {
|
|
return
|
|
}
|
|
|
|
viz := gocv.NewMat()
|
|
gocv.CvtColor(binary, &viz, gocv.ColorGrayToBGR)
|
|
defer viz.Close()
|
|
|
|
red := color.RGBA{R: 255, G: 0, B: 0, A: 255}
|
|
green := color.RGBA{R: 0, G: 255, B: 0, A: 255}
|
|
blue := color.RGBA{R: 0, G: 0, B: 255, A: 255}
|
|
yellow := color.RGBA{R: 255, G: 255, B: 0, A: 255}
|
|
|
|
// Draw all bands
|
|
for i, band := range bands {
|
|
bandColor := red
|
|
if i == tallestIdx {
|
|
bandColor = green // Graph band in green
|
|
} else if tallestIdx < len(bands)-1 && i == tallestIdx+1 {
|
|
bandColor = blue // Numbers band in blue
|
|
}
|
|
|
|
rect := image.Rect(0, band.minY, viz.Cols(), band.maxY)
|
|
gocv.Rectangle(&viz, rect, bandColor, 2)
|
|
|
|
label := fmt.Sprintf("Band #%d", i+1)
|
|
gocv.PutText(&viz, label, image.Pt(10, band.minY+20),
|
|
gocv.FontHersheyPlain, 1.5, yellow, 2)
|
|
}
|
|
|
|
gocv.IMWrite(filename, viz)
|
|
fmt.Printf("DEBUG: Saved visualization to %s\n", filename)
|
|
}
|
|
|
|
// VisualizeDigitBand highlights the digit band specifically
|
|
func VisualizeDigitBand(binary gocv.Mat, bands []Band, tallestIdx int, filename string) {
|
|
if !DEBUG || tallestIdx >= len(bands)-1 {
|
|
return
|
|
}
|
|
|
|
viz := gocv.NewMat()
|
|
gocv.CvtColor(binary, &viz, gocv.ColorGrayToBGR)
|
|
defer viz.Close()
|
|
|
|
red := color.RGBA{R: 255, G: 0, B: 0, A: 255}
|
|
green := color.RGBA{R: 0, G: 255, B: 0, A: 255}
|
|
blue := color.RGBA{R: 0, G: 0, B: 255, A: 255}
|
|
yellow := color.RGBA{R: 255, G: 255, B: 0, A: 255}
|
|
cyan := color.RGBA{R: 0, G: 255, B: 255, A: 255}
|
|
|
|
// Draw all bands with specific highlighting
|
|
for i, band := range bands {
|
|
bandColor := red
|
|
lineThickness := 1
|
|
labelText := fmt.Sprintf("Band #%d", i+1)
|
|
|
|
if i == tallestIdx {
|
|
bandColor = green
|
|
labelText = "GRAPH BAND"
|
|
} else if i == tallestIdx+1 {
|
|
// DIGIT BAND - highlight strongly
|
|
bandColor = cyan
|
|
lineThickness = 3
|
|
labelText = fmt.Sprintf("DIGIT BAND [Y: %d-%d]", band.minY, band.maxY)
|
|
}
|
|
|
|
rect := image.Rect(0, band.minY, viz.Cols(), band.maxY)
|
|
gocv.Rectangle(&viz, rect, bandColor, lineThickness)
|
|
|
|
// Add label
|
|
labelColor := yellow
|
|
if i == tallestIdx+1 {
|
|
labelColor = cyan
|
|
}
|
|
gocv.PutText(&viz, labelText, image.Pt(10, band.minY+20),
|
|
gocv.FontHersheyPlain, 1.5, labelColor, 2)
|
|
}
|
|
|
|
// Draw additional markers for digit band
|
|
if tallestIdx < len(bands)-1 {
|
|
digitBand := bands[tallestIdx+1]
|
|
|
|
// Draw vertical lines at digit band boundaries
|
|
for x := 50; x < viz.Cols(); x += 100 {
|
|
gocv.Line(&viz,
|
|
image.Pt(x, digitBand.minY),
|
|
image.Pt(x, digitBand.maxY),
|
|
blue, 1)
|
|
}
|
|
|
|
// Add dimension text
|
|
height := digitBand.maxY - digitBand.minY
|
|
dimText := fmt.Sprintf("Height: %dpx, Width: 860px", height)
|
|
gocv.PutText(&viz, dimText, image.Pt(10, digitBand.maxY+30),
|
|
gocv.FontHersheyPlain, 1.2, cyan, 2)
|
|
}
|
|
|
|
gocv.IMWrite(filename, viz)
|
|
fmt.Printf("DEBUG: Saved digit band visualization to %s\n", filename)
|
|
}
|
|
|
|
// ValidateBaselineTwoSegments checks if the baseline consists of two separate segments
|
|
// by counting black-to-white transitions along the baseline (should be exactly 1)
|
|
func ValidateBaselineTwoSegments(binary gocv.Mat, leftPoint, rightPoint *image.Point) bool {
|
|
// Calculate slope for Y interpolation along the baseline
|
|
deltaY := float64(rightPoint.Y - leftPoint.Y)
|
|
deltaX := float64(rightPoint.X - leftPoint.X)
|
|
slope := deltaY / deltaX
|
|
|
|
if DEBUG {
|
|
fmt.Printf("DEBUG: Validating baseline segments (slope=%.4f) from X=%d to X=%d\n",
|
|
slope, leftPoint.X, rightPoint.X)
|
|
}
|
|
|
|
// Count black-to-white transitions
|
|
transitions := 0
|
|
lastWasWhite := false
|
|
|
|
for x := leftPoint.X; x <= rightPoint.X; x++ {
|
|
// Interpolate Y position along the baseline, then go 2px up (away from thick edge)
|
|
y := leftPoint.Y + int(float64(x-leftPoint.X)*slope) - 2
|
|
isWhite := binary.GetUCharAt(y, x) == 255
|
|
|
|
// Detect black-to-white transition
|
|
if isWhite && !lastWasWhite {
|
|
transitions++
|
|
}
|
|
lastWasWhite = isWhite
|
|
}
|
|
|
|
if DEBUG {
|
|
fmt.Printf("DEBUG: Found %d black-to-white transitions\n", transitions)
|
|
}
|
|
|
|
// We need exactly 1 black-to-white transition for 2 segments
|
|
// (entering the second segment after the gap)
|
|
if transitions != 1 {
|
|
if DEBUG {
|
|
fmt.Printf("DEBUG: INVALID - Need exactly 1 transition, found %d\n", transitions)
|
|
}
|
|
return false
|
|
}
|
|
|
|
if DEBUG {
|
|
fmt.Println("DEBUG: VALID - Baseline has two distinct segments")
|
|
}
|
|
return true
|
|
}
|
|
|
|
// RotateImage rotates an image by the given angle in degrees
|
|
// The angle is the detected rotation that needs to be corrected
|
|
func RotateImage(img gocv.Mat, angleDegrees float64) gocv.Mat {
|
|
// Get image center
|
|
center := image.Point{
|
|
X: img.Cols() / 2,
|
|
Y: img.Rows() / 2,
|
|
}
|
|
|
|
// Apply rotation to correct the tilt
|
|
// NOTE: If image rotates wrong direction, change to: -angleDegrees
|
|
rotMat := gocv.GetRotationMatrix2D(center, angleDegrees, 1.0)
|
|
defer rotMat.Close()
|
|
|
|
// Apply rotation
|
|
rotated := gocv.NewMat()
|
|
gocv.WarpAffine(img, &rotated, rotMat, image.Point{X: img.Cols(), Y: img.Rows()})
|
|
|
|
if DEBUG {
|
|
fmt.Printf("DEBUG: Applied rotation: %.3f°\n", angleDegrees)
|
|
}
|
|
|
|
return rotated
|
|
}
|
|
|
|
// ScaleToWidth scales an image to the specified width while maintaining aspect ratio
|
|
func ScaleToWidth(img gocv.Mat, targetWidth int) gocv.Mat {
|
|
// Calculate scale factor
|
|
currentWidth := img.Cols()
|
|
scale := float64(targetWidth) / float64(currentWidth)
|
|
|
|
// Calculate new dimensions
|
|
newHeight := int(float64(img.Rows()) * scale)
|
|
|
|
// Resize image
|
|
scaled := gocv.NewMat()
|
|
gocv.Resize(img, &scaled, image.Point{X: targetWidth, Y: newHeight}, 0, 0, gocv.InterpolationLinear)
|
|
|
|
if DEBUG {
|
|
fmt.Printf("DEBUG: Scaled image from %dx%d to %dx%d (scale factor: %.3f)\n",
|
|
currentWidth, img.Rows(), targetWidth, newHeight, scale)
|
|
}
|
|
|
|
return scaled
|
|
}
|
|
|
|
// VisualizeSimpleBands draws just the graph and digit bands clearly
|
|
func VisualizeSimpleBands(binary gocv.Mat, graphBand, digitBand Band, filename string) {
|
|
if !DEBUG {
|
|
return
|
|
}
|
|
|
|
viz := gocv.NewMat()
|
|
gocv.CvtColor(binary, &viz, gocv.ColorGrayToBGR)
|
|
defer viz.Close()
|
|
|
|
green := color.RGBA{R: 0, G: 255, B: 0, A: 255}
|
|
cyan := color.RGBA{R: 0, G: 255, B: 255, A: 255}
|
|
|
|
// Draw graph band in green
|
|
graphRect := image.Rect(0, graphBand.minY, viz.Cols(), graphBand.maxY)
|
|
gocv.Rectangle(&viz, graphRect, green, 2)
|
|
gocv.PutText(&viz, "GRAPH BAND", image.Pt(10, graphBand.minY+20),
|
|
gocv.FontHersheyPlain, 1.5, green, 2)
|
|
|
|
// Draw digit band in cyan with thick border
|
|
digitRect := image.Rect(0, digitBand.minY, viz.Cols(), digitBand.maxY)
|
|
gocv.Rectangle(&viz, digitRect, cyan, 3)
|
|
label := fmt.Sprintf("DIGIT BAND [Y: %d-%d, H: %dpx]",
|
|
digitBand.minY, digitBand.maxY, digitBand.maxY-digitBand.minY)
|
|
gocv.PutText(&viz, label, image.Pt(10, digitBand.minY+20),
|
|
gocv.FontHersheyPlain, 1.5, cyan, 2)
|
|
|
|
gocv.IMWrite(filename, viz)
|
|
fmt.Printf("DEBUG: Saved 2-band visualization to %s\n", filename)
|
|
}
|
|
|
|
// ExtractBandRegion extracts a specific band region from an image
|
|
func ExtractBandRegion(img gocv.Mat, band Band) gocv.Mat {
|
|
// Create a region of interest (ROI) for the band
|
|
rect := image.Rect(0, band.minY, img.Cols(), band.maxY)
|
|
region := img.Region(rect)
|
|
|
|
// Clone the region to create an independent Mat
|
|
extracted := gocv.NewMat()
|
|
region.CopyTo(&extracted)
|
|
|
|
if DEBUG {
|
|
fmt.Printf("DEBUG: Extracted band region [%d-%d], size: %dx%d\n",
|
|
band.minY, band.maxY, extracted.Cols(), extracted.Rows())
|
|
}
|
|
|
|
return extracted
|
|
}
|
|
|
|
// ScaleByFactor scales an image by a specific scale factor
|
|
func ScaleByFactor(img gocv.Mat, scaleFactor float64) gocv.Mat {
|
|
// Calculate new dimensions
|
|
newWidth := int(float64(img.Cols()) * scaleFactor)
|
|
newHeight := int(float64(img.Rows()) * scaleFactor)
|
|
|
|
// Resize image
|
|
scaled := gocv.NewMat()
|
|
gocv.Resize(img, &scaled, image.Point{X: newWidth, Y: newHeight}, 0, 0, gocv.InterpolationLinear)
|
|
|
|
if DEBUG {
|
|
fmt.Printf("DEBUG: Scaled image from %dx%d to %dx%d (scale factor: %.3f)\n",
|
|
img.Cols(), img.Rows(), newWidth, newHeight, scaleFactor)
|
|
}
|
|
|
|
return scaled
|
|
}
|