From 1b047eae8cb8ea884fff5284acb0bacffb7457a6 Mon Sep 17 00:00:00 2001 From: clonejo Date: Sun, 7 Jul 2024 17:57:50 +0200 Subject: [PATCH 1/3] Allow masking off part of the frame when stitching I curse yee, shrubbery! --- cmd/trainbot/main.go | 22 ++++++++++++++++++---- internal/pkg/stitch/auto.go | 1 + internal/pkg/stitch/stitch.go | 19 +++++++++++++++---- 3 files changed, 34 insertions(+), 8 deletions(-) diff --git a/cmd/trainbot/main.go b/cmd/trainbot/main.go index ebe5712..6357c70 100644 --- a/cmd/trainbot/main.go +++ b/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() diff --git a/internal/pkg/stitch/auto.go b/internal/pkg/stitch/auto.go index 61eed02..e2dfcfd 100644 --- a/internal/pkg/stitch/auto.go +++ b/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 { diff --git a/internal/pkg/stitch/stitch.go b/internal/pkg/stitch/stitch.go index b5aac3d..f4fd757 100644 --- a/internal/pkg/stitch/stitch.go +++ b/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) From 88e36b9a4e33c6cb70c7bbf32b214391f93ee5ef Mon Sep 17 00:00:00 2001 From: clonejo Date: Sun, 4 May 2025 12:45:23 +0200 Subject: [PATCH 2/3] Document how to create mask in GIMP --- README.md | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/README.md b/README.md index c15a06a..c67f242 100644 --- a/README.md +++ b/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 From ad02a941101e85ddbda9095f0fe7697fbf3ce94d Mon Sep 17 00:00:00 2001 From: clonejo Date: Wed, 13 May 2026 01:08:04 +0200 Subject: [PATCH 3/3] Add mask preview to confighelper I still get displayed white in areas of partial transparency, but already usable. --- internal/pkg/server/wwwdata/index.html | 122 +++++++++++++++++++++---- 1 file changed, 106 insertions(+), 16 deletions(-) diff --git a/internal/pkg/server/wwwdata/index.html b/internal/pkg/server/wwwdata/index.html index 7c5be79..30ac676 100644 --- a/internal/pkg/server/wwwdata/index.html +++ b/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,7 +72,7 @@ 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) { @@ -68,9 +80,11 @@ } 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,7 +151,6 @@ 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} @@ -122,7 +158,7 @@ RECT_W=${rounded.w} RECT_H=${rounded.h}` - selectionState.active = false + selectionState.waitForSecondClick = false } renderState() @@ -130,7 +166,51 @@ 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,8 +228,14 @@ #rect { position: absolute; border: 1px solid red; + pointer-events: none; } + #rect img { + position: absolute; + pointer-events: none; + } + #grid { position: absolute; pointer-events: none; @@ -157,7 +243,6 @@ 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 @@

Trainbot Confighelper

Click on the video (twice) to select a rectangle.
-
+
+ + 

 	

+	

+		
+