Progress bars

cmd/mosaic/main.go

@@ -5,16 +5,16 @@ import (
 	"flag"
 	"fmt"
 	"image"
-	"log"
 	"os"
 	"os/signal"
 	"path/filepath"
 	"runtime"
 	"strings"
-	"time"
 
 	"github.com/pboyd/mosaic"
 
+	"github.com/schollz/progressbar/v3"
+
 	"image/gif"
 	"image/jpeg"
 	"image/png"
@@ -98,32 +98,39 @@ func main() {
 
 	index := mosaic.NewIndex(config)
 	index.StatusHandler = func(imgs <-chan mosaic.IndexImage) {
+		bar := progressbar.Default(-1)
 		for ii := range imgs {
 			if ii.Err != nil {
-				fmt.Fprintf(os.Stderr, "Error indexing %s: %v\n", ii.Path, ii.Err)
+				fmt.Fprintf(os.Stderr, "\rError indexing %s: %v\n", ii.Path, ii.Err)
 			} else {
-				log.Printf("%s: %.6x", ii.Path, ii.Color)
+				//nolint:errcheck
+				bar.Add(1)
 			}
 		}
 	}
 
-	start := time.Now()
 	err = index.AddPath(ctx, tileImagesPath)
-	elapsed := time.Since(start)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "Error indexing tile images: %v\n", err)
 		os.Exit(1)
 	}
 
-	fmt.Printf("Indexed %d images in %s\n", index.Len(), elapsed)
-
 	generator := mosaic.NewGenerator(config, index)
-	generator.StatusHandler = func(imgs <-chan mosaic.GeneratorStatus) {
+	generator.StatusHandler = func(total int64, imgs <-chan mosaic.GeneratorStatus) {
+		bar := progressbar.Default(total)
+		last := 0
 		for gs := range imgs {
 			if gs.Err != nil {
-				fmt.Fprintf(os.Stderr, "Error generating %s: %v\n", gs.Path, gs.Err)
-			} else {
-				log.Printf("%d/%d: %s", gs.TileNumber, gs.TotalTiles, gs.Path)
+				fmt.Fprintf(os.Stderr, "\rError generating %s: %v\n", gs.Path, gs.Err)
+				continue
+			}
+
+			// TileNumber may be out of order, so we need to keep track of the last
+			// tile number we saw.
+			if gs.TileNumber > last {
+				//nolint:errcheck
+				bar.Add(gs.TileNumber - last)
+				last = gs.TileNumber
 			}
 		}
 	}

colors.go

@@ -1,13 +1,16 @@
 package mosaic
 
 import (
+	"errors"
 	"image"
 	"image/color"
-	"math"
 
 	color_extractor "github.com/marekm4/color-extractor"
 )
 
+// ErrNoPrimaryColor is returned when no primary color is found in the image.
+var ErrNoPrimaryColor = errors.New("no primary color found")
+
 func colorVector(c uint32) []float64 {
 	return []float64{
 		float64(c >> 16),         // r
@@ -16,15 +19,15 @@ func colorVector(c uint32) []float64 {
 	}
 }
 
-func primaryColor(img image.Image, smallBucket float64) uint32 {
+func primaryColor(img image.Image, smallBucket float64) (uint32, error) {
 	colors := color_extractor.ExtractColorsWithConfig(img, color_extractor.Config{
 		SmallBucket: smallBucket,
 		DownSizeTo:  224,
 	})
 	if len(colors) == 0 {
-		return math.MaxUint32
+		return 0, ErrNoPrimaryColor
 	}
-	return colorRGB(colors[0])
+	return colorRGB(colors[0]), nil
 }
 
 func colorRGB(c color.Color) uint32 {

colors_test.go

@@ -15,7 +15,7 @@ func TestPrimaryColor(t *testing.T) {
 	crop := img.(interface {
 		SubImage(r image.Rectangle) image.Image
 	}).SubImage(image.Rect(10, 10, 20, 20))
-	c := primaryColor(crop, 0.01)
-
+	c, err := primaryColor(crop, 0.01)
+	assert.NoError(err)
 	assert.Equal(uint32(0xff00ff), c)
 }

go.mod

@@ -12,6 +12,12 @@ require (
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/mattn/go-runewidth v0.0.14 // indirect
+	github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/rivo/uniseg v0.2.0 // indirect
+	github.com/schollz/progressbar/v3 v3.13.1 // indirect
+	golang.org/x/sys v0.6.0 // indirect
+	golang.org/x/term v0.6.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -1,17 +1,35 @@
 github.com/cdipaolo/goml v0.0.0-20220715001353-00e0c845ae1c h1:uqJXOhayPfl/QruVBP6VF0KUWNDzO/F14X8CPEkkFD8=
 github.com/cdipaolo/goml v0.0.0-20220715001353-00e0c845ae1c/go.mod h1:Ue8jgVLdBDCtsh1laikvraXqXzKCyKiruCcCcaeNDFE=
+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/disintegration/imaging v1.6.2 h1:w1LecBlG2Lnp8B3jk5zSuNqd7b4DXhcjwek1ei82L+c=
 github.com/disintegration/imaging v1.6.2/go.mod h1:44/5580QXChDfwIclfc/PCwrr44amcmDAg8hxG0Ewe4=
+github.com/k0kubun/go-ansi v0.0.0-20180517002512-3bf9e2903213/go.mod h1:vNUNkEQ1e29fT/6vq2aBdFsgNPmy8qMdSay1npru+Sw=
 github.com/marekm4/color-extractor v1.2.1 h1:3Zb2tQsn6bITZ8MBVhc33Qn1k5/SEuZ18mrXGUqIwn0=
 github.com/marekm4/color-extractor v1.2.1/go.mod h1:90VjmiHI6M8ez9eYUaXLdcKnS+BAOp7w+NpwBdkJmpA=
+github.com/mattn/go-isatty v0.0.17/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
+github.com/mattn/go-runewidth v0.0.14 h1:+xnbZSEeDbOIg5/mE6JF0w6n9duR1l3/WmbinWVwUuU=
+github.com/mattn/go-runewidth v0.0.14/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
+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/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
+github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
+github.com/schollz/progressbar/v3 v3.13.1 h1:o8rySDYiQ59Mwzy2FELeHY5ZARXZTVJC7iHD6PEFUiE=
+github.com/schollz/progressbar/v3 v3.13.1/go.mod h1:xvrbki8kfT1fzWzBT/UZd9L6GA+jdL7HAgq2RFnO6fQ=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 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-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.6.0 h1:MVltZSvRTcU2ljQOhs94SXPftV6DCNnZViHeQps87pQ=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/term v0.6.0 h1:clScbb1cHjoCkyRbWwBEUZ5H/tIFu5TAXIqaZD0Gcjw=
+golang.org/x/term v0.6.0/go.mod h1:m6U89DPEgQRMq3DNkDClhWw02AUbt2daBVO4cn4Hv9U=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.6 h1:aRYxNxv6iGQlyVaZmk6ZgYEDa+Jg18DxebPSrd6bg1M=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=

index.go

@@ -183,7 +183,7 @@ func (idx *Index) processOne(path string) (uint32, error) {
 	if idx.config.ResizeTiles {
 		img = imaging.Fill(img, idx.config.TileWidth, idx.config.TileHeight, imaging.Center, imaging.Lanczos)
 	}
-	return primaryColor(img, idx.config.IndexThreshold), nil
+	return primaryColor(img, idx.config.IndexThreshold)
 }
 
 func (idx *Index) mergeColorChannels(chs ...<-chan IndexImage) <-chan IndexImage {

mosaic.go

@@ -3,7 +3,6 @@ package mosaic
 import (
 	"context"
 	"image"
-	"math"
 	"math/rand"
 	"os"
 	"sync"
@@ -33,7 +32,7 @@ type Config struct {
 type Generator struct {
 	config        Config
 	index         *Index
-	StatusHandler func(<-chan GeneratorStatus)
+	StatusHandler func(total int64, progress <-chan GeneratorStatus)
 }
 
 // NewGenerator creates a new Generator.
@@ -48,7 +47,6 @@ func NewGenerator(config Config, index *Index) *Generator {
 type GeneratorStatus struct {
 	Bounds     image.Rectangle
 	TileNumber int
-	TotalTiles int
 	Path       string
 	Err        error
 }
@@ -71,12 +69,12 @@ func (g *Generator) Generate(ctx context.Context, src image.Image) image.Image {
 		statusChans[i] = g.matchAndSwapTiles(output, tiles)
 	}
 
-	g.wait(statusChans)
+	g.wait(src.Bounds(), statusChans)
 
 	return output
 }
 
-func (*Generator) defaultStatusHandler(ch <-chan GeneratorStatus) {
+func (*Generator) defaultStatusHandler(total int64, ch <-chan GeneratorStatus) {
 	for range ch {
 		// do nothing
 	}
@@ -135,8 +133,10 @@ func (g *Generator) matchAndSwapTiles(output draw.Image, tiles <-chan image.Imag
 		for tile := range tiles {
 			status := tileStats(output.Bounds(), tile.Bounds())
 
-			c := primaryColor(tile, 0.01)
-			if c == math.MaxUint32 {
+			c, err := primaryColor(tile, 0.01)
+			if err != nil {
+				status.Err = err
+				out <- status
 				continue
 			}
 
@@ -169,19 +169,19 @@ func tileStats(imageBounds, tileBounds image.Rectangle) GeneratorStatus {
 	col := tileBounds.Min.Y / tileBounds.Dy()
 	return GeneratorStatus{
 		Bounds:     tileBounds,
-		TotalTiles: (imageBounds.Dx() / tileBounds.Dx()) * (imageBounds.Dy() / tileBounds.Dy()),
 		TileNumber: 1 + row + col*(imageBounds.Dx()/tileBounds.Dx()),
 	}
 }
 
-func (g *Generator) wait(statusChans []<-chan GeneratorStatus) {
+func (g *Generator) wait(bounds image.Rectangle, statusChans []<-chan GeneratorStatus) {
 	statusCh := make(chan GeneratorStatus, g.config.Workers*2)
 	statusHandler := g.defaultStatusHandler
 	if g.StatusHandler != nil {
 		statusHandler = g.StatusHandler
 	}
 
-	go statusHandler(statusCh)
+	totalTiles := int64((bounds.Dx() / g.config.TileWidth) * (bounds.Dy() / g.config.TileHeight))
+	go statusHandler(totalTiles, statusCh)
 
 	var wg sync.WaitGroup
 	wg.Add(len(statusChans))