Masking

README.md

@@ -264,6 +264,33 @@ libcamera-vid \
 mkvmerge -o test.mkv --timecodes 0:vid-timestamps.txt vid.h264
 ```
 
+## How to create mask
+
+- fetch a background image
+  1. either from confighelper
+      - select your rectangle in confighelper
+      - configure rectangle in trainbot
+      - take a screenshot, scale must be 1:1
+      - open screenshot in GIMP
+      - select rectangle (dimensions must be exact!)
+      - Image -> Crop to Selection
+  2. or from gif (video clip)
+      - open train .gif in gimp
+      - in the layers list, delete layers from the top until you find a good one
+      - flatten image
+- in layer list, right click layer, add mask
+- select the mask thumbnail
+- Using paintbrush tool, hide the obstructions (foreground) in black, white will be kept. Prefer to leave only a vertical pillar visible, to avoid artifacts.
+- CTRL-A, CTRL-C
+- create new file
+- CTRL-V, merge / anchor floating layer down
+- Colors -> Color to Alpha
+- Color: black, OK
+- File -> Export to mask.png
+- configure `RECT_MASK=mask.png`, upload mask.png to where you run trainbot
+
+If you end up with black artifacts, the pillar you let through is too small for the speed this train is travelling at.
+
 ## Code notes
 
 * Zerolog is used as logging framework

cmd/trainbot/main.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"image"
+	"image/png"
 	"io/fs"
 	"math"
 	"math/rand"
@@ -37,10 +38,11 @@ type config struct {
 	CameraW            int    `arg:"--camera-w,env:CAMERA_W" default:"1920" help:"Camera frame size width, ignored if using video file or picam3" placeholder:"X"`
 	CameraH            int    `arg:"--camera-h,env:CAMERA_H" default:"1080" help:"Camera frame size height, ignored if using video file or picam3" placeholder:"Y"`
 
-	RectX uint `arg:"-X,--rect-x,env:RECT_X" help:"Rect to look at, x (left)" placeholder:"N"`
-	RectY uint `arg:"-Y,--rect-y,env:RECT_Y" help:"Rect to look at, y (top)" placeholder:"N"`
-	RectW uint `arg:"-W,--rect-w,env:RECT_W" help:"Rect to look at, width" placeholder:"N"`
-	RectH uint `arg:"-H,--rect-h,env:RECT_H" help:"Rect to look at, height" placeholder:"N"`
+	RectX    uint    `arg:"-X,--rect-x,env:RECT_X" help:"Rect to look at, x (left)" placeholder:"N"`
+	RectY    uint    `arg:"-Y,--rect-y,env:RECT_Y" help:"Rect to look at, y (top)" placeholder:"N"`
+	RectW    uint    `arg:"-W,--rect-w,env:RECT_W" help:"Rect to look at, width" placeholder:"N"`
+	RectH    uint    `arg:"-H,--rect-h,env:RECT_H" help:"Rect to look at, height" placeholder:"N"`
+	RectMask *string `arg:"--mask,env:RECT_MASK" help:"When stitching, only take pixels from the white areas in the mask." placeholder:"FILE"`
 
 	Rotate180 bool `arg:"--rotate-180,env:ROTATE_180" help:"Rotate camera picture 180 degrees (only picam3)"`
 
@@ -152,12 +154,24 @@ func detectTrainsForever(c config, trainsOut chan<- *stitch.Train) {
 	defer src.Close()
 	srcBuf := vid.NewSrcBuf(src, failedFramesMax)
 
+	var mask image.Image
+	if c.RectMask != nil {
+		fMask, err := os.Open(*c.RectMask)
+		if err != nil {
+			log.Panic().Err(err)
+		}
+		mask, err = png.Decode(fMask)
+		if err != nil {
+			log.Panic().Err(err)
+		}
+	}
 	stitcher := stitch.NewAutoStitcher(stitch.Config{
 		PixelsPerM:          c.PixelsPerM,
 		MinSpeedKPH:         c.MinSpeedKPH,
 		MaxSpeedKPH:         c.MaxSpeedKPH,
 		MinLengthM:          c.MinLengthM,
 		MaxFrameCountPerSeq: c.MaxFrameCountPerSeq,
+		Mask:                mask,
 	})
 	defer func() {
 		train := stitcher.TryStitchAndReset()

internal/pkg/server/wwwdata/index.html

@@ -29,16 +29,28 @@
 			})
 		}
 
-		function elementToImgCoordinates(event) {
-			const stream = event.target;
+		function elementToImgCoordinates(x, y) {
+			const stream = document.querySelector("#stream");
 			return {
-				x: stream.naturalWidth / stream.width * event.layerX,
-				y: stream.naturalHeight / stream.height * event.layerY
+				x: stream.naturalWidth / stream.width * x,
+				y: stream.naturalHeight / stream.height * y,
+			}
+		}
+		function imgToElementCoordinates(x, y) {
+			const stream = document.querySelector("#stream");
+			return {
+				x: stream.width / stream.naturalWidth * x,
+				y: stream.height / stream.naturalHeight * y,
 			}
 		}
 
 		const selectionState = {
-			active: false,
+			waitForSecondClick: false,
+			mask: false,
+			maskSize: {
+				w: -1,
+				h: -1,
+			},
 			rect: {
 				x0: -1,
 				y0: -1,
@@ -60,17 +72,19 @@
 			rect.style.width = (selectionState.rect.x1 - selectionState.rect.x0)+"px"
 			rect.style.height = (selectionState.rect.y1 - selectionState.rect.y0)+"px"
 
-			rect.hidden = selectionState.rect.x0 < 0 || selectionState.active
+			rect.hidden = selectionState.rect.x0 < 0 || selectionState.waitForSecondClick
 		}
 
 		function swap(obj, key1, key2) {
 			[obj[key1], obj[key2]] = [obj[key2], obj[key1]]
 		}
 
 		function updateStreamRect(event) {
-			const imgCoord = elementToImgCoordinates(event)
+			const imgCoord = elementToImgCoordinates(event.layerX, event.layerY)
+			var rectCompleted = false
 
-			if (!selectionState.active) {
+			if (!selectionState.waitForSecondClick) {
+				// first click
 				selectionState.rect.x0 = event.layerX
 				selectionState.rect.y0 = event.layerY
 				selectionState.rect.x1 = -1
@@ -81,10 +95,24 @@
 				selectionState.img.x1 = -1
 				selectionState.img.y1 = -1
 
-				document.querySelector("#rectOptions").innerHTML = ""
-				document.querySelector("#rectOptionsEnv").innerHTML = ""
-				selectionState.active = true
-			} else {
+				rectCompleted = false
+			}
+
+			console.log(selectionState)
+
+			if (!selectionState.waitForSecondClick && selectionState.mask) {
+				// With a mask configured, we take width and height from the mask instead of letting the user click again.
+				wh = imgToElementCoordinates(selectionState.maskSize.w, selectionState.maskSize.h)
+				selectionState.rect.x1 = event.layerX + wh.x
+				selectionState.rect.y1 = event.layerY + wh.y
+				selectionState.img.x1 = imgCoord.x + selectionState.maskSize.w
+				selectionState.img.y1 = imgCoord.y + selectionState.maskSize.h
+
+				rectCompleted = true
+			}
+
+			if (selectionState.waitForSecondClick) {
+				// second click
 				selectionState.rect.x1 = event.layerX
 				selectionState.rect.y1 = event.layerY
 
@@ -101,6 +129,15 @@
 					swap(selectionState.img, "y0", "y1")
 				}
 
+				rectCompleted = true
+			}
+
+			if (!rectCompleted) {
+				document.querySelector("#rectOptions").innerHTML = ""
+				document.querySelector("#rectOptionsEnv").innerHTML = ""
+
+				selectionState.waitForSecondClick = true
+			} else {
 				const [x, y, w, h] = [
 					selectionState.img.x0,
 					selectionState.img.y0,
@@ -114,23 +151,66 @@
 					w: Math.round(w),
 					h: Math.round(h),
 				}
-
 				document.querySelector("#rectOptions").innerHTML = `-X ${rounded.x} -Y ${rounded.y} -W ${rounded.w} -H ${rounded.h}`
 				document.querySelector("#rectOptionsEnv").innerHTML =
 `RECT_X=${rounded.x}
 RECT_Y=${rounded.y}
 RECT_W=${rounded.w}
 RECT_H=${rounded.h}`
 
-				selectionState.active = false
+				selectionState.waitForSecondClick = false
 			}
 
 			renderState()
 			event.preventDefault()
 			return false
 		}
 
+		async function addMask(event) {
+			const [file] = document.querySelector("#maskFile").files;
+			if (file) {
+				const img = await invertMask(file);
+				console.log(img)
+				//document.querySelector("#rect").style.maskImage = `url("${url}")`
+				document.querySelector("#mask").src = img.url;
+				selectionState.mask = true;
+				selectionState.maskSize = {w: img.w, h: img.h};
+			}
+		}
+
+		async function invertMask(file) {
+			const url = URL.createObjectURL(file);
+
+			const img = await new Promise((res, rej) => {
+				const i = new Image();
+				i.onload = () => res(i);
+				i.onerror = rej;
+				i.src = url;
+			});
+			URL.revokeObjectURL(url);
+
+			const canvas = document.createElement('canvas');
+			canvas.width = img.width;
+			canvas.height = img.height;
+			const ctx = canvas.getContext('2d');
+			ctx.drawImage(img, 0, 0);
+
+			const imageData = ctx.getImageData(0, 0, canvas.width, canvas.height);
+			for (let i = 0; i < imageData.data.length; i += 4) {
+				// turn to black:
+				imageData.data[0] = 200;
+				imageData.data[1] = 0;
+				imageData.data[2] = 0;
+				// invert alpha:
+				imageData.data[i+3] = 255 - imageData.data[i+3];
+			}
+			ctx.putImageData(imageData, 0, 0);
+
+			return new Promise(res => canvas.toBlob(blob => res({url: URL.createObjectURL(blob), w: img.width, h: img.height})));
+		}
+
 		window.onload = function(){
+			document.querySelector("#previewMask").addEventListener("click", addMask)
 			document.querySelector("#detectCameras").addEventListener("click", detectCamerasClick)
 			document.querySelector("#stream").addEventListener("click", updateStreamRect)
 
@@ -148,16 +228,21 @@
 		#rect {
 			position: absolute;
 			border: 1px solid red;
+			pointer-events: none;
 		}
 
+		#rect img {
+			position: absolute;
+			pointer-events: none;
+    }
+
 		#grid {
 			position: absolute;
 			pointer-events: none;
 			top: 0;
 			left: 0;
 			width: 100%;
 			height: 100%;
-
 			opacity: 0.8;
 			background-image:  linear-gradient(#3366ff 1px, transparent 1px), linear-gradient(to right, #3366ff 1px, transparent 1px);
 			background-size: 50px 50px;
@@ -173,12 +258,17 @@
 	<h1>Trainbot Confighelper</h1>
 	<div>Click on the video (twice) to select a rectangle.</div>
 	<div id="container">
-		<div id="rect"></div>
+		<div id="rect">
+			<img id="mask">
+		</div>
 		<div id="grid"></div>
 		<img id="stream" src="stream.mjpeg" />
 	</div>
 	<pre id="rectOptions"></pre>
 	<pre id="rectOptionsEnv"></pre>
+	<p>
+		<input type="file" name="mask" id="maskFile"><button id="previewMask">Preview mask</button>
+	</p>
 	<button id="detectCameras">Detect v4l cameras</button>
 	<pre id="cameraOptions"></pre>
 </body>

internal/pkg/stitch/auto.go

@@ -29,6 +29,7 @@ type Config struct {
 	MaxSpeedKPH         float64
 	MinLengthM          float64
 	MaxFrameCountPerSeq int
+	Mask                image.Image
 }
 
 func (c *Config) minPxPerFrame(framePeriodS float64) int {

internal/pkg/stitch/stitch.go

@@ -39,7 +39,7 @@ func sign(x float64) float64 {
 	return 0
 }
 
-func stitch(frames []image.Image, dx []int) (*image.RGBA, error) {
+func stitch(frames []image.Image, dx []int, mask image.Image) (*image.RGBA, error) {
 	t0 := time.Now()
 	defer func() {
 		log.Trace().Dur("dur", time.Since(t0)).Msg("stitch() duration")
@@ -60,6 +60,9 @@ func stitch(frames []image.Image, dx []int) (*image.RGBA, error) {
 			log.Panic().Msg("frame bounds or size not consistent, this should not happen")
 		}
 	}
+	if mask != nil && fb != mask.Bounds() {
+		log.Panic().Interface("frame", fb).Interface("mask", mask.Bounds()).Msg("mask size and frame size do not match!")
+	}
 
 	// Calculate base width.
 	sign := isign(dx[0])
@@ -79,18 +82,26 @@ func stitch(frames []image.Image, dx []int) (*image.RGBA, error) {
 	}
 	img := image.NewRGBA(rect)
 
+	mp := image.Point{}
+	op := draw.Src
+	if mask != nil {
+		log.Debug().Msg("Stitching using mask.")
+		mp = mask.Bounds().Min
+		op = draw.Over
+	}
+
 	// Forward?
 	if w > 0 {
 		pos := 0
 		for i, f := range frames {
-			draw.Draw(img, img.Bounds().Add(image.Pt(pos, 0)), f, f.Bounds().Min, draw.Src)
+			draw.DrawMask(img, img.Bounds().Add(image.Pt(pos, 0)), f, f.Bounds().Min, mask, mp, op)
 			pos += dx[i]
 		}
 	} else {
 		// Backwards.
 		pos := -w - fb.Dx()
 		for i, f := range frames {
-			draw.Draw(img, img.Bounds().Add(image.Pt(pos, 0)), f, f.Bounds().Min, draw.Src)
+			draw.DrawMask(img, img.Bounds().Add(image.Pt(pos, 0)), f, f.Bounds().Min, mask, mp, op)
 			pos += dx[i]
 		}
 	}
@@ -235,7 +246,7 @@ func fitAndStitch(seq sequence, c Config) (*Train, error) {
 		return nil, fmt.Errorf("discarded because too slow, %f < %f", speed, c.minSpeedPxPS())
 	}
 
-	img, err := stitch(seq.frames, dxFit)
+	img, err := stitch(seq.frames, dxFit, c.Mask)
 	if err != nil {
 		prometheus.RecordFitAndStitchResult("unable_to_assemble_image")
 		return nil, fmt.Errorf("unable to assemble image: %w", err)