Update stitching program

- Divide workload into chunks - Add multithreading - Add console messages - Add progress bar - Unload tile images after some time - Fix median filter
2024-11-18 17:17:31 +00:00 · 2019-10-24 00:28:22 +02:00 · 2019-10-24 00:28:22 +02:00 · b11ba98db7
commit b11ba98db7
parent a228b00815
7 changed files with 177 additions and 60 deletions
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@ -1,8 +1,11 @@
 {
    "cSpell.words": [
        "Vogel",
+        "gridify",
        "hacky",
        "kbinani",
-        "noita"
+        "mapcap",
+        "noita",
+        "schollz"
    ]
 }
--- a/bin/stitch/imagetile.go
+++ b/bin/stitch/imagetile.go
@ -10,19 +10,26 @@ import (
 	"image"
 	_ "image/png"
 	"os"
+	"sync"
+	"time"
 )

 type imageTile struct {
 	fileName string

 	offset image.Point // Correction offset of the image, so that it aligns pixel perfect with other images. Determined by image matching.
-	image  image.Image // Either a rectangle or an RGBA image. The bounds of this image are determined by the filename.
+
+	image      image.Image // Either a rectangle or an RGBA image. The bounds of this image are determined by the filename.
+	imageMutex *sync.Mutex
 }

-func (it *imageTile) loadImage() error {
+func (it *imageTile) GetImage() (*image.RGBA, error) {
+	it.imageMutex.Lock()
+	defer it.imageMutex.Unlock() // TODO: Use RWMutex
+
 	// Check if the image is already loaded
-	if _, ok := it.image.(*image.RGBA); ok {
-		return nil
+	if img, ok := it.image.(*image.RGBA); ok {
+		return img, nil
 	}

 	// Store rectangle of the old image
@ -30,18 +37,18 @@ func (it *imageTile) loadImage() error {

 	file, err := os.Open(it.fileName)
 	if err != nil {
-		return err
+		return &image.RGBA{}, err
 	}
 	defer file.Close()

 	img, _, err := image.Decode(file)
 	if err != nil {
-		return err
+		return &image.RGBA{}, err
 	}

 	imgRGBA, ok := img.(*image.RGBA)
 	if !ok {
-		return fmt.Errorf("Expected an RGBA image, got %T instead", img)
+		return &image.RGBA{}, fmt.Errorf("Expected an RGBA image, got %T instead", img)
 	}

 	// Restore the position of the image rectangle
@ -49,11 +56,30 @@ func (it *imageTile) loadImage() error {

 	it.image = imgRGBA

-	return nil
+	// Free the image after some time
+	go func() {
+		time.Sleep(5 * time.Second)
+
+		it.imageMutex.Lock()
+		defer it.imageMutex.Unlock()
+		it.image = it.image.Bounds()
+	}()
+
+	return imgRGBA, nil
 }

-func (it *imageTile) unloadImage() {
-	it.image = it.image.Bounds()
+func (it *imageTile) OffsetBounds() image.Rectangle {
+	it.imageMutex.Lock()
+	defer it.imageMutex.Unlock() // TODO: Use RWMutex
+
+	return it.image.Bounds().Add(it.offset)
+}
+
+func (it *imageTile) Bounds() image.Rectangle {
+	it.imageMutex.Lock()
+	defer it.imageMutex.Unlock() // TODO: Use RWMutex
+
+	return it.image.Bounds()
 }

 func (it *imageTile) String() string {
--- a/bin/stitch/imagetiles.go
+++ b/bin/stitch/imagetiles.go
@ -14,8 +14,12 @@ import (
 	"math/rand"
 	"path/filepath"
 	"regexp"
+	"runtime"
 	"sort"
 	"strconv"
+	"sync"
+
+	"github.com/schollz/progressbar/v2"
 )

 const tileAlignmentSearchRadius = 5
@ -62,27 +66,27 @@ func loadImages(path string) ([]imageTile, error) {
 		}

 		imageTiles = append(imageTiles, imageTile{
-			fileName: file,
-			image:    image.Rect(x, y, x+width, y+height),
+			fileName:   file,
+			image:      image.Rect(x, y, x+width, y+height),
+			imageMutex: &sync.Mutex{},
 		})
 	}

 	return imageTiles, nil
 }

-// alignTilePair returns the pixel delta for the first tile, so that it aligns perfectly with the second.
+// AlignTilePair returns the pixel delta for the first tile, so that it aligns perfectly with the second.
 // This function will load images if needed.
-func alignTilePair(tileA, tileB *imageTile, searchRadius int) (image.Point, error) {
-	if err := tileA.loadImage(); err != nil {
+func AlignTilePair(tileA, tileB *imageTile, searchRadius int) (image.Point, error) {
+	imgA, err := tileA.GetImage()
+	if err != nil {
 		return image.Point{}, err
 	}
-	if err := tileB.loadImage(); err != nil {
+	imgB, err := tileB.GetImage()
+	if err != nil {
 		return image.Point{}, err
 	}

-	// Type assertion.
-	imgA, imgB := *tileA.image.(*image.RGBA), *tileB.image.(*image.RGBA)
-
 	bestPoint := image.Point{}
 	bestValue := math.Inf(1)

@ -90,7 +94,7 @@ func alignTilePair(tileA, tileB *imageTile, searchRadius int) (image.Point, erro
 		for x := -searchRadius; x <= searchRadius; x++ {
 			point := image.Point{x, y} // Offset of the first image.

-			value := getImageDifferenceValue(&imgA, &imgB, point)
+			value := getImageDifferenceValue(imgA, imgB, point)
 			if bestValue > value {
 				bestValue, bestPoint = value, point
 			}
@ -100,7 +104,7 @@ func alignTilePair(tileA, tileB *imageTile, searchRadius int) (image.Point, erro
 	return bestPoint, nil
 }

-func (tp tilePairs) alignTiles(tiles []*imageTile) error {
+func (tp tilePairs) AlignTiles(tiles []*imageTile) error {

 	n := len(tiles)
 	maxOperations, operations := (n-1)*(n)/2, 0
@ -113,7 +117,7 @@ func (tp tilePairs) alignTiles(tiles []*imageTile) error {
 			_, ok := tp[tileAlignmentKeys{tileA, tileB}]
 			if !ok {
 				// Entry doesn't exist yet. Determine tile pair alignment.
-				offset, err := alignTilePair(tileA, tileB, tileAlignmentSearchRadius)
+				offset, err := AlignTilePair(tileA, tileB, tileAlignmentSearchRadius)
 				if err != nil {
 					return fmt.Errorf("Failed to align tile pair %v %v: %w", tileA, tileB, err)
 				}
@ -200,13 +204,13 @@ func (tp tilePairs) alignTiles(tiles []*imageTile) error {
 	return nil
 }

-func (tp tilePairs) stitch(tiles []imageTile, destImage *image.RGBA) error {
+func (tp tilePairs) Stitch(tiles []imageTile, destImage *image.RGBA) error {
 	intersectTiles := []*imageTile{}

 	// Get only the tiles that intersect with the destination image bounds.
 	// Ignore alignment here, doesn't matter if an image overlaps a few pixels anyways.
 	for i, tile := range tiles {
-		if tile.image.Bounds().Add(tile.offset).Overlaps(destImage.Bounds()) {
+		if tile.OffsetBounds().Overlaps(destImage.Bounds()) {
 			intersectTiles = append(intersectTiles, &tiles[i])
 		}
 	}
@ -225,44 +229,79 @@ func (tp tilePairs) stitch(tiles []imageTile, destImage *image.RGBA) error {
 		draw.Draw(destImage, destImage.Bounds(), intersectTile.image, destImage.Bounds().Min, draw.Over)
 	}*/

-	drawMedianBlended(intersectTiles, destImage)
-
 	/*for _, intersectTile := range intersectTiles {
 		drawLabel(destImage, intersectTile.image.Bounds().Min.X, intersectTile.image.Bounds().Min.Y, fmt.Sprintf("%v", intersectTile.fileName))
 	}*/

-	return nil
+	return drawMedianBlended(intersectTiles, destImage)
 }

-func drawMedianBlended(tiles []*imageTile, destImage *image.RGBA) {
-	bounds := destImage.Bounds()
+// StitchGrid calls stitch, but divides the workload into a grid of chunks.
+// Additionally it runs the workload multithreaded.
+func (tp tilePairs) StitchGrid(tiles []imageTile, destImage *image.RGBA, gridSize int) (errResult error) {
+	workloads := gridifyRectangle(destImage.Bounds(), gridSize)

-	// Make sure images are loaded.
-	for _, tile := range tiles {
-		tile.loadImage()
+	bar := progressbar.New(len(workloads))
+	bar.RenderBlank()
+
+	// Start worker threads
+	wc := make(chan image.Rectangle)
+	wg := sync.WaitGroup{}
+	for i := 0; i < runtime.NumCPU(); i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for workload := range wc {
+				if err := tp.Stitch(tiles, destImage.SubImage(workload).(*image.RGBA)); err != nil {
+					errResult = err // This will not stop execution, but at least one of any errors is returned.
+				}
+				bar.Add(1)
+			}
+		}()
 	}

+	// Push workload to worker threads
+	for _, workload := range workloads {
+		wc <- workload
+	}
+
+	// Wait until all worker threads are done
+	close(wc)
+	wg.Wait()
+
+	// Newline because of the progress bar
+	fmt.Println("")
+
+	return
+}
+
+func drawMedianBlended(tiles []*imageTile, destImage *image.RGBA) error {
+	bounds := destImage.Bounds()
+
 	for iy := bounds.Min.Y; iy < bounds.Max.Y; iy++ {
 		for ix := bounds.Min.X; ix < bounds.Max.X; ix++ {
-			//colList := []color.RGBA{}
-			rList, gList, bList := []int{}, []int{}, []int{}
+			rList, gList, bList := []int16{}, []int16{}, []int16{}
 			point := image.Point{ix, iy}
+			found := false

 			// Iterate through all tiles, and create a list of colors.
 			for _, tile := range tiles {
 				tilePoint := point.Sub(tile.offset)
-				imageRGBA, ok := tile.image.(*image.RGBA)
-				if ok && tilePoint.In(imageRGBA.Bounds()) {
+				imageRGBA, err := tile.GetImage() // TODO: Optimize, as it's slow to get tiles and images every pixel
+				if err != nil {
+					return fmt.Errorf("Couldn't load image: %w", err)
+				}
+				if tilePoint.In(imageRGBA.Bounds().Inset(4)) { // Reduce image bounds by 4 pixels on each side, because otherwise there will be artifacts.
 					col := imageRGBA.RGBAAt(tilePoint.X, tilePoint.Y)
-					//colList = append(colList, col)
-					rList, gList, bList = append(rList, int(col.R)), append(gList, int(col.G)), append(bList, int(col.B))
+					rList, gList, bList = append(rList, int16(col.R)), append(gList, int16(col.G)), append(bList, int16(col.B))
+					found = true
 				}
 			}

-			// Sort color list.
-			/*sort.Slice(colList, func(i, j int) bool {
-				return rgbToHSV(colList[i]) < rgbToHSV(colList[j])
-			})*/
+			// If there were no tiles to get data from, ignore the pixel.
+			if !found {
+				continue
+			}

 			// Sort rList.
 			sort.Slice(rList, func(i, j int) bool {
@ -279,36 +318,26 @@ func drawMedianBlended(tiles []*imageTile, destImage *image.RGBA) {
 				return bList[i] < bList[j]
 			})

-			//var col color.RGBA
-
-			/*if len(colList)%2 == 0 {
-				// Even
-				a, b := colList[len(colList)/2-1], colList[len(colList)/2]
-				col = color.RGBA{uint8((uint16(a.R) + uint16(b.R)) / 2), uint8((uint16(a.G) + uint16(b.G)) / 2), uint8((uint16(a.B) + uint16(b.B)) / 2), uint8((uint16(a.A) + uint16(b.A)) / 2)}
-			} else if len(colList) > 0 {
-				// Odd
-				col = colList[(len(colList)-1)/2]
-			}*/
-
+			// Take the middle element of each color.
 			var r, g, b uint8
 			if len(rList)%2 == 0 {
 				// Even
 				r = uint8((rList[len(rList)/2-1] + rList[len(rList)/2]) / 2)
-			} else if len(rList) > 0 {
+			} else {
 				// Odd
 				r = uint8(rList[(len(rList)-1)/2])
 			}
 			if len(gList)%2 == 0 {
 				// Even
 				g = uint8((gList[len(gList)/2-1] + gList[len(gList)/2]) / 2)
-			} else if len(gList) > 0 {
+			} else {
 				// Odd
 				g = uint8(gList[(len(gList)-1)/2])
 			}
 			if len(bList)%2 == 0 {
 				// Even
 				b = uint8((bList[len(bList)/2-1] + bList[len(bList)/2]) / 2)
-			} else if len(bList) > 0 {
+			} else {
 				// Odd
 				b = uint8(bList[(len(bList)-1)/2])
 			}
@ -316,4 +345,6 @@ func drawMedianBlended(tiles []*imageTile, destImage *image.RGBA) {
 			destImage.SetRGBA(ix, iy, color.RGBA{r, g, b, 255})
 		}
 	}
+
+	return nil
 }
--- a/bin/stitch/stitch.go
+++ b/bin/stitch/stitch.go
@ -16,10 +16,12 @@ import (
 var inputPath = filepath.Join(".", "..", "..", "output")

 func main() {
+	log.Printf("Starting to read tile information at \"%v\"", inputPath)
 	tiles, err := loadImages(inputPath)
 	if err != nil {
 		log.Panic(err)
 	}
+	log.Printf("Got %v tiles", len(tiles))

 	/*f, err := os.Create("cpu.prof")
 	if err != nil {
@ -31,13 +33,18 @@ func main() {
 	}
 	defer pprof.StopCPUProfile()*/

-	outputImage := image.NewRGBA(image.Rect(-4000, -4000, 8000, 8000))
+	outputRect := image.Rect(-10000, -10000, 10000, 10000)

+	log.Printf("Creating output image with a size of %v", outputRect.Size())
+	outputImage := image.NewRGBA(outputRect)
+
+	log.Printf("Stitching %v tiles into an image at %v", len(tiles), outputImage.Bounds())
 	tp := make(tilePairs)
-	if err := tp.stitch(tiles, outputImage); err != nil {
+	if err := tp.StitchGrid(tiles, outputImage, 256); err != nil {
 		log.Panic(err)
 	}

+	log.Printf("Creating output file \"%v\"", "output.png")
 	f, err := os.Create("output.png")
 	if err != nil {
 		log.Panic(err)
@ -51,4 +58,6 @@ func main() {
 	if err := f.Close(); err != nil {
 		log.Panic(err)
 	}
+	log.Printf("Created output file \"%v\"", "output.png")
+
 }
--- a/bin/stitch/util.go
+++ b/bin/stitch/util.go
@ -62,9 +62,22 @@ func getImageDifferenceValue(a, b *image.RGBA, offsetA image.Point) float64 {
 	return float64(value) / float64(intersectionWidth*intersectionHeight)
 }

+func gridifyRectangle(rect image.Rectangle, gridSize int) (result []image.Rectangle) {
+	for y := divideFloor(rect.Min.Y, gridSize); y < divideCeil(rect.Max.Y, gridSize); y++ {
+		for x := divideFloor(rect.Min.X, gridSize); x < divideCeil(rect.Max.X, gridSize); x++ {
+			tempRect := image.Rect(x*gridSize, y*gridSize, (x+1)*gridSize, (y+1)*gridSize)
+			if tempRect.Overlaps(rect) {
+				result = append(result, tempRect)
+			}
+		}
+	}
+
+	return
+}
+
 func drawLabel(img *image.RGBA, x, y int, label string) {
 	col := color.RGBA{200, 100, 0, 255}
-	point := fixed.Point26_6{fixed.Int26_6(x * 64), fixed.Int26_6(y * 64)}
+	point := fixed.Point26_6{X: fixed.Int26_6(x * 64), Y: fixed.Int26_6(y * 64)}

 	d := &font.Drawer{
 		Dst:  img,
@ -91,3 +104,25 @@ func pointAbs(p image.Point) image.Point {
 	}
 	return p
 }
+
+// Integer division that rounds to the next integer towards negative infinity
+func divideFloor(a, b int) int {
+	temp := a / b
+
+	if ((a ^ b) < 0) && (a%b != 0) {
+		return temp - 1
+	}
+
+	return temp
+}
+
+// Integer division that rounds to the next integer towards positive infinity
+func divideCeil(a, b int) int {
+	temp := a / b
+
+	if ((a ^ b) >= 0) && (a%b != 0) {
+		return temp + 1
+	}
+
+	return temp
+}
--- a/go.mod
+++ b/go.mod
@ -7,5 +7,6 @@ require (
 	github.com/gen2brain/shm v0.0.0-20180314170312-6c18ff7f8b90 // indirect
 	github.com/kbinani/screenshot v0.0.0-20190719135742-f06580e30cdc
 	github.com/lxn/win v0.0.0-20190919090605-24c5960b03d8 // indirect
+	github.com/schollz/progressbar/v2 v2.14.0
 	golang.org/x/image v0.0.0-20191009234506-e7c1f5e7dbb8
 )
--- a/go.sum
+++ b/go.sum
@ -1,11 +1,23 @@
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802 h1:1BDTz0u9nC3//pOCMdNH+CiXJVYJh5UQNCOBG7jbELc=
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/gen2brain/shm v0.0.0-20180314170312-6c18ff7f8b90 h1:QagTG5rauLt6pVVEhnVSrlIX4ifhVIZOwmw6x6D8TUw=
 github.com/gen2brain/shm v0.0.0-20180314170312-6c18ff7f8b90/go.mod h1:uF6rMu/1nvu+5DpiRLwusA6xB8zlkNoGzKn8lmYONUo=
 github.com/kbinani/screenshot v0.0.0-20190719135742-f06580e30cdc h1:kGFotla6Dyr6a2ILeExAHlttPgJtnoP/GIw2uVN/4h4=
 github.com/kbinani/screenshot v0.0.0-20190719135742-f06580e30cdc/go.mod h1:f8GY5V3lRzakvEyr49P7hHRYoHtPr8zvj/7JodCoRzw=
 github.com/lxn/win v0.0.0-20190919090605-24c5960b03d8 h1:RVMGIuuNgrpGB7I79f6xfhGCkpN47IaEGh8VTM0p7Xc=
 github.com/lxn/win v0.0.0-20190919090605-24c5960b03d8/go.mod h1:ouWl4wViUNh8tPSIwxTVMuS014WakR1hqvBc2I0bMoA=
+github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db h1:62I3jR2EmQ4l5rM/4FEfDWcRD+abF5XlKShorW5LRoQ=
+github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db/go.mod h1:l0dey0ia/Uv7NcFFVbCLtqEBQbrT4OCwCSKTEv6enCw=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/schollz/progressbar/v2 v2.14.0 h1:vo7bdkI9E4/CIk9DnL5uVIaybLQiVtiCC2vO+u9j5IM=
+github.com/schollz/progressbar/v2 v2.14.0/go.mod h1:6YZjqdthH6SCZKv2rqGryrxPtfmRB/DWZxSMfCXPyD8=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 golang.org/x/image v0.0.0-20191009234506-e7c1f5e7dbb8 h1:hVwzHzIUGRjiF7EcUjqNxk3NCfkPxbDKRdnNE1Rpg0U=
 golang.org/x/image v0.0.0-20191009234506-e7c1f5e7dbb8/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0=
 golang.org/x/sys v0.0.0-20190904154756-749cb33beabd h1:DBH9mDw0zluJT/R+nGuV3jWFWLFaHyYZWD4tOT+cjn0=