feat: Add masking support for improved OCR

This allows masking artifacts in the image that could be confusing for
the OCR process, like units or unrelated values on the meter screen.

Author: dcelasun

README.md

@@ -123,6 +123,8 @@ All options can be set via CLI flags or environment variables. Flags take preced
 | `--crop` | `CROP` | *(disabled)* | Crop rectangle as `x0,y0,x1,y1` applied before OCR |
 | `--ocr-match-regex` | `OCR_MATCH_REGEX` | `^000\d+$` | Regex to identify the meter reading from OCR text |
 | `--ocr-fix-regex` | `OCR_FIX_REGEX` | *(disabled)* | Comma-separated list of regex substitutions as `pattern=replacement` applied in order before matching (e.g. `^O=0,^030=000`) |
+| `--ocr-mask-regions` | `OCR_MASK_REGIONS` | *(disabled)* | Comma-separated rectangle coordinates to mask before OCR, as `x1,y1,x2,y2[,x3,y3,x4,y4,...]` (applied after crop) |
+| `--ocr-mask-colors` | `OCR_MASK_COLORS` | `000000` | Comma-separated hex colors for mask regions. One color applies to all; otherwise must match the number of regions |
 | `--mqtt-broker` | `MQTT_BROKER` | *(disabled)* | MQTT broker URL, e.g. `tcp://192.168.1.100:1883` |
 | `--mqtt-user` | `MQTT_USER` | | MQTT username |
 | `--mqtt-password` | `MQTT_PASSWORD` | | MQTT password |

http.go

@@ -60,20 +60,30 @@ func handleOCR(w http.ResponseWriter, r *http.Request) {
 }
 
 func processOCR(imageData []byte, batLevel, batVoltage int) {
-	var croppedData []byte
+	var cropped, masked bool
 	ocrData := imageData
 	if cropRect != nil {
-		cropped, err := cropImage(imageData, cropRect)
+		data, err := cropImage(imageData, cropRect)
 		if err != nil {
 			log.Printf("crop error: %v, using original image", err)
 		} else {
-			croppedData = cropped
-			ocrData = cropped
+			ocrData = data
+			cropped = true
+		}
+	}
+
+	if len(ocrMasks) > 0 {
+		data, err := maskImage(ocrData, ocrMasks)
+		if err != nil {
+			log.Printf("mask error: %v, using unmasked image", err)
+		} else {
+			ocrData = data
+			masked = true
 		}
 	}
 
 	// Store images to disk before OCR so we have them even if OCR fails.
-	imagePath := storeImages(imageData, croppedData)
+	imagePath := storeImages(imageData, ocrData, cropped, masked)
 
 	tmpDir, err := os.MkdirTemp("", "ocr-")
 	if err != nil {

main.go

@@ -20,6 +20,7 @@ var (
 	cropRect    *cropConfig
 	ocrMatchRe  *regexp.Regexp
 	ocrFixRules []ocrFixRule
+	ocrMasks    []maskRegion
 
 	mqttBroker             string
 	mqttUser               string
@@ -76,6 +77,16 @@ func main() {
 				Usage:   "Comma-separated list of regex substitutions applied to OCR text before matching, each as pattern=replacement (e.g. ^O=0,^030=000)",
 				Sources: cli.EnvVars("OCR_FIX_REGEX"),
 			},
+			&cli.StringFlag{
+				Name:    "ocr-mask-regions",
+				Usage:   "Comma-separated rectangle coordinates to mask before OCR, as x1,y1,x2,y2[,x3,y3,x4,y4,...] (applied after crop)",
+				Sources: cli.EnvVars("OCR_MASK_REGIONS"),
+			},
+			&cli.StringFlag{
+				Name:    "ocr-mask-colors",
+				Usage:   "Comma-separated hex colors for mask regions (e.g. 0099CC). One color applies to all regions; otherwise must match the number of regions",
+				Sources: cli.EnvVars("OCR_MASK_COLORS"),
+			},
 			&cli.StringFlag{
 				Name:    "mqtt-broker",
 				Usage:   "MQTT broker URL, e.g. tcp://192.168.1.100:1883 (disabled if empty)",
@@ -143,6 +154,7 @@ func run(_ context.Context, cmd *cli.Command) error {
 	cropRect = parseCropRect(cmd.String("crop"))
 	ocrMatchRe = regexp.MustCompile(cmd.String("ocr-match-regex"))
 	ocrFixRules = parseOCRFixRules(cmd.String("ocr-fix-regex"))
+	ocrMasks = parseMaskRegions(cmd.String("ocr-mask-regions"), cmd.String("ocr-mask-colors"))
 
 	mqttBroker = cmd.String("mqtt-broker")
 	mqttUser = cmd.String("mqtt-user")

ocr.go

@@ -2,9 +2,12 @@ package main
 
 import (
 	"bytes"
+	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"image"
+	"image/color"
+	"image/draw"
 	"image/jpeg"
 	"log"
 	"os/exec"
@@ -76,6 +79,122 @@ func cropImage(data []byte, c *cropConfig) ([]byte, error) {
 	return buf.Bytes(), nil
 }
 
+type maskRegion struct {
+	Rect  image.Rectangle
+	Color color.RGBA
+}
+
+func parseMaskRegions(regions, colors string) []maskRegion {
+	regions = strings.TrimSpace(regions)
+	if regions == "" {
+		return nil
+	}
+
+	coords := strings.Split(regions, ",")
+	if len(coords)%4 != 0 {
+		log.Fatalf("ocr-mask-regions: must have groups of 4 coordinates (x1,y1,x2,y2), got %d values", len(coords))
+	}
+
+	vals := make([]int, len(coords))
+	for i, c := range coords {
+		v, err := strconv.Atoi(strings.TrimSpace(c))
+		if err != nil {
+			log.Fatalf("ocr-mask-regions: invalid integer %q: %v", c, err)
+		}
+		vals[i] = v
+	}
+
+	numRects := len(vals) / 4
+	rectColors := parseMaskColors(colors, numRects)
+
+	masks := make([]maskRegion, numRects)
+	for i := 0; i < numRects; i++ {
+		x0, y0 := vals[i*4], vals[i*4+1]
+		x1, y1 := vals[i*4+2], vals[i*4+3]
+		if x0 >= x1 || y0 >= y1 {
+			log.Fatalf("ocr-mask-regions: invalid rectangle #%d (x0,y0 must be < x1,y1): %d,%d,%d,%d", i+1, x0, y0, x1, y1)
+		}
+		masks[i] = maskRegion{
+			Rect:  image.Rect(x0, y0, x1, y1),
+			Color: rectColors[i],
+		}
+	}
+	return masks
+}
+
+func parseMaskColors(s string, numRects int) []color.RGBA {
+	s = strings.TrimSpace(s)
+	defaultColor := color.RGBA{0, 0, 0, 255}
+
+	if s == "" {
+		colors := make([]color.RGBA, numRects)
+		for i := range colors {
+			colors[i] = defaultColor
+		}
+		return colors
+	}
+
+	parts := strings.Split(s, ",")
+	parsed := make([]color.RGBA, 0, len(parts))
+	for _, p := range parts {
+		p = strings.TrimSpace(p)
+		if p == "" {
+			continue
+		}
+		parsed = append(parsed, parseHexColor(p))
+	}
+
+	if len(parsed) == 1 {
+		colors := make([]color.RGBA, numRects)
+		for i := range colors {
+			colors[i] = parsed[0]
+		}
+		return colors
+	}
+
+	if len(parsed) != numRects {
+		log.Fatalf("ocr-mask-colors: must specify 1 color or exactly %d (one per rectangle), got %d", numRects, len(parsed))
+	}
+	return parsed
+}
+
+func parseHexColor(s string) color.RGBA {
+	s = strings.TrimPrefix(s, "#")
+	if len(s) != 6 {
+		log.Fatalf("ocr-mask-colors: invalid hex color %q (must be 6 hex digits like 0099CC)", s)
+	}
+	b, err := hex.DecodeString(s)
+	if err != nil {
+		log.Fatalf("ocr-mask-colors: invalid hex color %q: %v", s, err)
+	}
+	return color.RGBA{R: b[0], G: b[1], B: b[2], A: 255}
+}
+
+func maskImage(data []byte, masks []maskRegion) ([]byte, error) {
+	img, _, err := image.Decode(bytes.NewReader(data))
+	if err != nil {
+		return nil, fmt.Errorf("decode image: %w", err)
+	}
+
+	bounds := img.Bounds()
+	dst := image.NewRGBA(bounds)
+	draw.Draw(dst, bounds, img, bounds.Min, draw.Src)
+
+	for _, m := range masks {
+		r := m.Rect.Intersect(bounds)
+		if r.Empty() {
+			continue
+		}
+		draw.Draw(dst, r, &image.Uniform{m.Color}, image.Point{}, draw.Src)
+	}
+
+	var buf bytes.Buffer
+	if err := jpeg.Encode(&buf, dst, &jpeg.Options{Quality: 95}); err != nil {
+		return nil, fmt.Errorf("encode masked image: %w", err)
+	}
+	return buf.Bytes(), nil
+}
+
 type ocrOutput struct {
 	Texts  []string  `json:"texts"`
 	Scores []float64 `json:"scores"`