diff --git a/README.md b/README.md
index c15a06ae..c67f2426 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
diff --git a/cmd/trainbot/main.go b/cmd/trainbot/main.go
index ebe57123..6357c700 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/server/wwwdata/index.html b/internal/pkg/server/wwwdata/index.html
index 7c5be794..30ac676f 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.
-
+
+
![]()
+
+
+
diff --git a/internal/pkg/stitch/auto.go b/internal/pkg/stitch/auto.go
index 61eed022..e2dfcfd6 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 b5aac3d4..f4fd7572 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)