Automate the Boring Stuff with Go, Ch. 18: Images

// open an image file, say what it is, and read one pixel
package main

import (
	"fmt"
	"image"
	_ "image/jpeg" // registering decoders is how image.Decode learns a format
	_ "image/png"
	"os"
)

func main() {
	if len(os.Args) != 2 {
		fmt.Fprintln(os.Stderr, "usage: image-info <file>")
		os.Exit(2)
	}
	f, err := os.Open(os.Args[1])
	if err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
	defer f.Close()

	img, format, err := image.Decode(f)
	if err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
	b := img.Bounds()
	fmt.Printf("format: %s\nsize:   %d x %d\n", format, b.Dx(), b.Dy())

	// colours come back as 16-bit channels; shift by 8 for the familiar 0-255
	r, g, bl, a := img.At(b.Dx()/2, b.Dy()/2).RGBA()
	fmt.Printf("centre pixel: rgba(%d, %d, %d, %d)\n", r>>8, g>>8, bl>>8, a>>8)
}

// make an image from nothing: a background, a rectangle, a circle, saved twice
package main

import (
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"image/jpeg"
	"image/png"
	"os"
)

func main() {
	img := image.NewRGBA(image.Rect(0, 0, 400, 300))

	// fill the whole thing with one colour
	draw.Draw(img, img.Bounds(), image.NewUniform(color.RGBA{245, 240, 225, 255}), image.Point{}, draw.Src)

	// a rectangle is a Draw onto a sub-rectangle
	draw.Draw(img, image.Rect(40, 40, 200, 160), image.NewUniform(color.RGBA{200, 60, 40, 255}), image.Point{}, draw.Src)

	// a circle is a loop over pixels inside its bounding box
	cx, cy, r := 290, 170, 80
	for y := cy - r; y <= cy+r; y++ {
		for x := cx - r; x <= cx+r; x++ {
			if (x-cx)*(x-cx)+(y-cy)*(y-cy) <= r*r {
				img.Set(x, y, color.RGBA{40, 90, 200, 255})
			}
		}
	}

	if err := save("shapes.png", func(f *os.File) error { return png.Encode(f, img) }); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
	if err := save("shapes.jpg", func(f *os.File) error { return jpeg.Encode(f, img, &jpeg.Options{Quality: 90}) }); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
	fmt.Println("wrote shapes.png and shapes.jpg")
}

func save(name string, encode func(*os.File) error) error {
	f, err := os.Create(name)
	if err != nil {
		return err
	}
	if err := encode(f); err != nil {
		f.Close()
		return err
	}
	return f.Close()
}

// shrink an image to half size with a good resampling filter
package main

import (
	"fmt"
	"image"
	_ "image/jpeg"
	"image/png"
	"os"

	"golang.org/x/image/draw"
)

func main() {
	if len(os.Args) != 3 {
		fmt.Fprintln(os.Stderr, "usage: resize <in> <out.png>")
		os.Exit(2)
	}
	in, err := os.Open(os.Args[1])
	if err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
	defer in.Close()
	src, _, err := image.Decode(in)
	if err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}

	b := src.Bounds()
	dst := image.NewRGBA(image.Rect(0, 0, b.Dx()/2, b.Dy()/2))
	// CatmullRom is slow and sharp; ApproxBiLinear is fast and soft. Thumbnails want sharp.
	draw.CatmullRom.Scale(dst, dst.Bounds(), src, b, draw.Over, nil)

	out, err := os.Create(os.Args[2])
	if err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
	defer out.Close()
	if err := png.Encode(out, dst); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
	fmt.Printf("%dx%d -> %dx%d, wrote %s\n", b.Dx(), b.Dy(), dst.Bounds().Dx(), dst.Bounds().Dy(), os.Args[2])
}

// thumbnails makes a small copy of every image in a folder, longest side -size pixels,
// into a thumbs/ sub-folder. Resampling is CPU work, so the files are processed by a
// pool of workers (chapter 14).
//
//	go run ./ch18/projects/thumbnails -size 300 ./ch18/projects/thumbnails/sample
package main

import (
	"errors"
	"flag"
	"fmt"
	"image"
	"image/jpeg"
	_ "image/png"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"golang.org/x/image/draw"
	"golang.org/x/sync/errgroup"
)

// Fit returns the size that fits w x h inside a max x max box, keeping the aspect ratio.
// Images already smaller than the box are left alone.
func Fit(w, h, max int) (int, int) {
	if w <= max && h <= max {
		return w, h
	}
	if w >= h {
		return max, h * max / w
	}
	return w * max / h, max
}

// Thumbnail returns src scaled to fit in a max x max box.
func Thumbnail(src image.Image, max int) image.Image {
	b := src.Bounds()
	w, h := Fit(b.Dx(), b.Dy(), max)
	dst := image.NewRGBA(image.Rect(0, 0, w, h))
	draw.CatmullRom.Scale(dst, dst.Bounds(), src, b, draw.Src, nil)
	return dst
}

// One reads src, writes a JPEG thumbnail to dst, and returns the thumbnail's size.
func One(src, dst string, max int) (int, int, error) {
	in, err := os.Open(src)
	if err != nil {
		return 0, 0, err
	}
	defer in.Close()
	img, _, err := image.Decode(in)
	if err != nil {
		return 0, 0, fmt.Errorf("decode %s: %w", src, err)
	}
	thumb := Thumbnail(img, max)
	out, err := os.Create(dst)
	if err != nil {
		return 0, 0, err
	}
	if err := jpeg.Encode(out, thumb, &jpeg.Options{Quality: 85}); err != nil {
		out.Close()
		return 0, 0, err
	}
	b := thumb.Bounds()
	return b.Dx(), b.Dy(), out.Close()
}

// Process thumbnails every .png, .jpg and .jpeg directly inside folder into
// folder/thumbs, with at most workers running at once. It returns how many it wrote.
func Process(folder string, max, workers int) (int, error) {
	entries, err := os.ReadDir(folder)
	if err != nil {
		return 0, err
	}
	outDir := filepath.Join(folder, "thumbs")
	if err := os.MkdirAll(outDir, 0o755); err != nil {
		return 0, err
	}
	var g errgroup.Group
	g.SetLimit(workers)
	done := make(chan string, len(entries))
	for _, e := range entries {
		ext := strings.ToLower(filepath.Ext(e.Name()))
		if e.IsDir() || (ext != ".png" && ext != ".jpg" && ext != ".jpeg") {
			continue
		}
		src := filepath.Join(folder, e.Name())
		dst := filepath.Join(outDir, strings.TrimSuffix(e.Name(), filepath.Ext(e.Name()))+".jpg")
		g.Go(func() error {
			w, h, err := One(src, dst, max)
			if err != nil {
				return err
			}
			done <- fmt.Sprintf("%s -> %dx%d", e.Name(), w, h)
			return nil
		})
	}
	err = g.Wait()
	close(done)
	n := 0
	for line := range done {
		fmt.Println(line)
		n++
	}
	return n, err
}

func run(args []string) error {
	flags := flag.NewFlagSet("thumbnails", flag.ContinueOnError)
	size := flags.Int("size", 300, "longest side of a thumbnail, in pixels")
	workers := flags.Int("workers", runtime.NumCPU(), "images to process at once")
	if err := flags.Parse(args); err != nil {
		return err
	}
	if flags.NArg() != 1 {
		return errors.New("usage: thumbnails [-size 300] [-workers N] <folder>")
	}
	n, err := Process(flags.Arg(0), *size, *workers)
	if err != nil {
		return err
	}
	fmt.Printf("%d thumbnail(s) in %s\n", n, filepath.Join(flags.Arg(0), "thumbs"))
	return nil
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

// watermark-images stamps text, a logo, or both onto the bottom-right corner of every
// image in a folder, writing copies to a watermarked/ sub-folder. The originals are not
// touched.
//
//	go run ./ch18/projects/watermark-images -text "© Toul Cranmer" -logo ./ch18/projects/watermark-images/sample/logo.png ./ch18/projects/watermark-images/sample
package main

import (
	"errors"
	"flag"
	"fmt"
	"image"
	"image/color"
	"image/draw"
	_ "image/jpeg"
	"image/png"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/image/font"
	"golang.org/x/image/font/gofont/goregular"
	"golang.org/x/image/font/opentype"
	"golang.org/x/image/math/fixed"
)

// LoadFace builds a font face from the Go font that ships with x/image, so there is no
// font file to find on the reader's machine.
func LoadFace(size float64) (font.Face, error) {
	f, err := opentype.Parse(goregular.TTF)
	if err != nil {
		return nil, err
	}
	return opentype.NewFace(f, &opentype.FaceOptions{Size: size, DPI: 72})
}

// StampText draws text in the bottom-right corner: a dark shadow, then light text on top,
// so it reads on both bright and dark photos.
func StampText(dst *image.RGBA, text string, face font.Face, margin int) {
	d := &font.Drawer{Dst: dst, Face: face}
	width := d.MeasureString(text).Ceil()
	b := dst.Bounds()
	x := b.Max.X - margin - width
	y := b.Max.Y - margin
	d.Src = image.NewUniform(color.RGBA{0, 0, 0, 160})
	d.Dot = fixed.P(x+2, y+2)
	d.DrawString(text)
	d.Src = image.NewUniform(color.RGBA{255, 255, 255, 220})
	d.Dot = fixed.P(x, y)
	d.DrawString(text)
}

// Overlay composites logo onto the bottom-right corner, above any text, respecting the
// logo's own transparency.
func Overlay(dst *image.RGBA, logo image.Image, margin, liftBy int) {
	lb := logo.Bounds()
	b := dst.Bounds()
	at := image.Rect(b.Max.X-margin-lb.Dx(), b.Max.Y-margin-liftBy-lb.Dy(), b.Max.X-margin, b.Max.Y-margin-liftBy)
	draw.Draw(dst, at, logo, lb.Min, draw.Over)
}

// Watermark returns a stamped copy of src. Either text or logo may be empty/nil.
func Watermark(src image.Image, text string, face font.Face, logo image.Image) *image.RGBA {
	b := src.Bounds()
	dst := image.NewRGBA(image.Rect(0, 0, b.Dx(), b.Dy()))
	draw.Draw(dst, dst.Bounds(), src, b.Min, draw.Src)
	margin := max(12, b.Dx()/50)
	lift := 0
	if text != "" {
		StampText(dst, text, face, margin)
		lift = face.Metrics().Height.Ceil() + margin/2
	}
	if logo != nil {
		Overlay(dst, logo, margin, lift)
	}
	return dst
}

func decode(path string) (image.Image, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	img, _, err := image.Decode(f)
	if err != nil {
		return nil, fmt.Errorf("decode %s: %w", path, err)
	}
	return img, nil
}

func run(args []string) error {
	flags := flag.NewFlagSet("watermark-images", flag.ContinueOnError)
	text := flags.String("text", "", "text to stamp")
	logoPath := flags.String("logo", "", "PNG logo to stamp")
	fontSize := flags.Float64("font-size", 28, "text size in points")
	if err := flags.Parse(args); err != nil {
		return err
	}
	if flags.NArg() != 1 || (*text == "" && *logoPath == "") {
		return errors.New("usage: watermark-images [-text T] [-logo file.png] <folder>")
	}
	folder := flags.Arg(0)
	face, err := LoadFace(*fontSize)
	if err != nil {
		return err
	}
	var logo image.Image
	if *logoPath != "" {
		if logo, err = decode(*logoPath); err != nil {
			return err
		}
	}
	outDir := filepath.Join(folder, "watermarked")
	if err := os.MkdirAll(outDir, 0o755); err != nil {
		return err
	}
	entries, err := os.ReadDir(folder)
	if err != nil {
		return err
	}
	n := 0
	for _, e := range entries {
		ext := strings.ToLower(filepath.Ext(e.Name()))
		if e.IsDir() || (ext != ".png" && ext != ".jpg" && ext != ".jpeg") {
			continue
		}
		src := filepath.Join(folder, e.Name())
		if abs, _ := filepath.Abs(src); *logoPath != "" && abs == mustAbs(*logoPath) {
			continue // do not stamp the logo onto itself
		}
		img, err := decode(src)
		if err != nil {
			return err
		}
		out := filepath.Join(outDir, strings.TrimSuffix(e.Name(), ext)+".png")
		f, err := os.Create(out)
		if err != nil {
			return err
		}
		if err := png.Encode(f, Watermark(img, *text, face, logo)); err != nil {
			f.Close()
			return err
		}
		if err := f.Close(); err != nil {
			return err
		}
		fmt.Println("wrote", out)
		n++
	}
	fmt.Printf("%d image(s) watermarked\n", n)
	return nil
}

func mustAbs(p string) string {
	abs, err := filepath.Abs(p)
	if err != nil {
		return p
	}
	return abs
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}