Moji: A small Golang project to make an animated emoji mosaic from a source image

main.go

package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"image"
	"image/color"
	"image/color/palette"
	"image/gif"
	"image/jpeg"
	"image/png"
	"io/fs"
	"io/ioutil"
	"log"
	"math"
	"math/rand"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"golang.org/x/image/draw"

	"github.com/dhconnelly/rtreego"
)

var (
	CachePath string
)

const (
	// Where emoji images can be found. Relative to bin.
	EmojiPath = "emojis"
	// Input size of emoji images.
	EmojiSize = 72
	// How many nearest emojis to pick from randomly.
	EmojiJitter = 3
)

func init() {
	CachePath = filepath.Join("cache", "emoji.json")

	image.RegisterFormat("png", "png", png.Decode, png.DecodeConfig)
	image.RegisterFormat("jpg", "jpg", jpeg.Decode, jpeg.DecodeConfig)

	rand.Seed(time.Now().UnixNano())
}

type EmojiItem struct {
	name   string
	avgClr color.RGBA
}

func NewEmojiItem(name string, avgClr color.RGBA) *EmojiItem {
	return &EmojiItem{
		name:   name,
		avgClr: avgClr,
	}
}

func (e *EmojiItem) Name() string {
	return e.name
}

func (e *EmojiItem) Bounds() *rtreego.Rect {
	point := colorToPoint(e.avgClr)
	rect, err := rtreego.NewRectFromPoints(point, point)
	if err != nil {
		log.Fatal(err)
	}

	return rect
}

func main() {
	fmt.Println("Running...")

	args := os.Args[1:]
	inImgPath, outImgPath, pixelSize, scale, frames, err := parseArgs(args)
	if err != nil {
		log.Print(err)
		printUsage()
		os.Exit(1)
	}

	fmt.Printf("Loading emoji images for size: %d...\n", pixelSize)
	emojiImgs, err := emojiImages(EmojiPath, pixelSize)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println("Loading emoji colors...")
	emojiClrs, err := readOrCreateEmojiClrCache(emojiImgs, CachePath)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println("Creating search tree...")
	tree := createSearchTree(emojiClrs)

	fmt.Println("Loading input image...")
	img, err := loadImage(inImgPath)
	if err != nil {
		log.Fatal(err)
	}
	imgBounds := img.Bounds()

	scaledPX := int(math.Round(float64(imgBounds.Min.X) * scale))
	scaledPY := int(math.Round(float64(imgBounds.Min.Y) * scale))
	scaledQX := int(math.Round(float64(imgBounds.Max.X) * scale))
	scaledQY := int(math.Round(float64(imgBounds.Max.Y) * scale))
	scaledImgBounds := image.Rect(scaledPX, scaledPY, scaledQX, scaledQY)

	outImgs := make([]*image.Paletted, frames)
	delays := make([]int, frames)
	for i := 0; i < frames; i++ {
		outImgs[i] = image.NewPaletted(scaledImgBounds, palette.WebSafe)
		delays[i] = 50

		fmt.Printf("Writing output image frame: %d...\n", i)
		for y := scaledImgBounds.Min.Y; y < scaledImgBounds.Max.Y+pixelSize; y += pixelSize {
			for x := scaledImgBounds.Min.X; x < scaledImgBounds.Max.X+pixelSize; x += pixelSize {
				samplePX := int(math.Round(float64(x) / scale))
				samplePY := int(math.Round(float64(y) / scale))
				sampleQX := int(math.Round(float64(x+pixelSize) / scale))
				sampleQY := int(math.Round(float64(y+pixelSize) / scale))
				sampleRegion := image.Rect(samplePX, samplePY, sampleQX, sampleQY)

				dstPX := x
				dstPY := y
				dstQX := x + pixelSize
				dstQY := y + pixelSize
				dstRegion := image.Rect(dstPX, dstPY, dstQX, dstQY)

				pixels := imagePixels(img, sampleRegion)
				clr := averageColor(pixels)

				if clr.A == 0 {
					continue
				}

				emojiImg, err := findNearestEmoji(tree, emojiImgs, clr)
				if err != nil {
					log.Fatal(err)
				}

				draw.Draw(outImgs[i], dstRegion, emojiImg, image.Point{}, draw.Src)
			}
		}
	}

	f, err := os.Create(outImgPath)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	gif.EncodeAll(f, &gif.GIF{
		Image: outImgs,
		Delay: delays,
	})

	fmt.Println("Done.")
}

func printUsage() {
	log.Println("Usage: moji <input_image.(png|jpg)> <output_image.gif> <pixel_size [1, 72]> <scale (0,]> <frames (0,]>")
}

func parseArgs(args []string) (string, string, int, float64, int, error) {
	if len(args) != 5 {
		return "", "", 0, 0, 0, errors.New("not enough args")
	}

	inImgPath := args[0]
	outImgPath := args[1]
	pixelSizeStr := args[2]
	pixelSize, err := strconv.ParseInt(pixelSizeStr, 10, 0)
	if err != nil {
		return "", "", 0, 0, 0, err
	}
	if pixelSize <= 0 || pixelSize > EmojiSize {
		return "", "", 0, 0, 0, errors.New("invalid pixel size")
	}

	scaleStr := args[3]
	scale, err := strconv.ParseFloat(scaleStr, 64)
	if err != nil {
		return "", "", 0, 0, 0, err
	}
	if scale <= 0 {
		return "", "", 0, 0, 0, errors.New("invalid scale")
	}

	framesStr := args[4]
	frames, err := strconv.ParseInt(framesStr, 10, 0)
	if err != nil {
		return "", "", 0, 0, 0, err
	}
	if frames <= 0 {
		return "", "", 0, 0, 0, errors.New("invalid frames")
	}

	return inImgPath, outImgPath, int(pixelSize), scale, int(frames), nil
}

func loadImage(path string) (image.Image, error) {
	file, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	img, _, err := image.Decode(bytes.NewReader(file))
	if err != nil {
		return nil, err
	}

	return img, nil
}

func createSearchTree(emojiClrs map[string]color.RGBA) *rtreego.Rtree {
	rt := rtreego.NewTree(4, 25, 50)
	for name, clr := range emojiClrs {
		item := NewEmojiItem(name, clr)
		rt.Insert(item)
	}

	return rt
}

func findNearestEmoji(tree *rtreego.Rtree, emojiImgs map[string]image.Image, clr color.RGBA) (image.Image, error) {
	items := tree.NearestNeighbors(EmojiJitter, colorToPoint(clr))
	idx := rand.Intn(len(items))
	item, ok := items[idx].(*EmojiItem)
	if !ok {
		return nil, errors.New("tree item couldn't be cast to an EmojiItem")
	}

	name := item.Name()
	img := emojiImgs[name]

	return img, nil
}

func readOrCreateEmojiClrCache(emojiImgs map[string]image.Image, cachePath string) (map[string]color.RGBA, error) {
	var emojiClrs map[string]color.RGBA
	bytes, err := ioutil.ReadFile(CachePath)
	if err != nil {

		emojiClrs, err = createCache(emojiImgs, cachePath)
		if err != nil {
			return nil, err
		}
	} else {
		emojiClrs, err = parseCacheBytes(bytes)
		if err != nil {
			return nil, err
		}
	}

	return emojiClrs, nil
}

func parseCacheBytes(bytes []byte) (map[string]color.RGBA, error) {
	emojiClrs := map[string]color.RGBA{}

	err := json.Unmarshal(bytes, &emojiClrs)
	if err != nil {
		return nil, err
	}

	return emojiClrs, nil
}

func createCache(emojiImgs map[string]image.Image, cachePath string) (map[string]color.RGBA, error) {
	emojiClrs := emojiColors(emojiImgs)

	bytes, err := json.Marshal(emojiClrs)
	if err != nil {
		return nil, err
	}

	ioutil.WriteFile(cachePath, bytes, fs.ModePerm)

	return emojiClrs, nil
}

func emojiImages(assetPath string, size int) (map[string]image.Image, error) {
	res := map[string]image.Image{}

	files, err := ioutil.ReadDir(assetPath)
	if err != nil {
		return nil, err
	}

	for _, file := range files {
		name := file.Name()

		srcImg, err := loadImage(filepath.Join(assetPath, file.Name()))
		if err != nil {
			return nil, err
		}

		img := image.NewRGBA(image.Rect(0, 0, size, size))
		draw.NearestNeighbor.Scale(img, img.Bounds(), srcImg, srcImg.Bounds(), draw.Src, nil)

		res[name] = img
	}

	return res, nil
}

func emojiColors(emojiImgs map[string]image.Image) map[string]color.RGBA {
	res := map[string]color.RGBA{}

	for name, img := range emojiImgs {
		pxls := imagePixels(img, img.Bounds())
		avgClr := averageColor(pxls)

		res[name] = avgClr
	}

	return res
}

func imagePixels(img image.Image, region image.Rectangle) []color.RGBA {
	pX := region.Min.X
	pY := region.Min.Y
	qX := region.Max.X
	qY := region.Max.Y

	// If the provided region is outside of the image bounds, clamp it to
	// the nearest complete region within the image.
	imgBounds := img.Bounds()
	imgWidth := imgBounds.Max.X
	imgHeight := imgBounds.Max.Y
	if imgWidth < qX {
		pX = imgWidth - region.Dx()
	}
	if imgHeight < qY {
		pY = imgHeight - region.Dy()
	}

	pixels := []color.RGBA{}
	for y := pY; y < qY; y++ {
		for x := pX; x < qX; x++ {
			pixel := img.At(x, y)

			r, g, b, a := rgba32ToRGBA8(pixel.RGBA())
			rgba := color.RGBA{R: r, G: g, B: b, A: a}

			pixels = append(pixels, rgba)
		}
	}

	return pixels
}

func rgba32ToRGBA8(r32, g32, b32, a32 uint32) (uint8, uint8, uint8, uint8) {
	r := uint8(r32 >> 8)
	g := uint8(g32 >> 8)
	b := uint8(b32 >> 8)
	a := uint8(a32 >> 8)

	return r, g, b, a
}

func averageColor(clrs []color.RGBA) color.RGBA {
	rComp := uint64(0)
	gComp := uint64(0)
	bComp := uint64(0)
	aComp := uint64(0)
	pixelsCount := uint64(len(clrs))

	if pixelsCount == 0 {
		return color.RGBA{R: 0, G: 0, B: 0, A: 0}
	}

	for _, clr := range clrs {
		r, g, b, a := rgba32ToRGBA8(clr.RGBA())

		r64 := uint64(r)
		g64 := uint64(g)
		b64 := uint64(b)
		a64 := uint64(a)

		rComp += r64 * r64
		gComp += g64 * g64
		bComp += b64 * b64
		aComp += a64 * a64
	}

	rAvg := rComp / pixelsCount
	gAvg := gComp / pixelsCount
	bAvg := bComp / pixelsCount
	aAvg := aComp / pixelsCount

	r := uint8(math.Sqrt(float64(rAvg)))
	g := uint8(math.Sqrt(float64(gAvg)))
	b := uint8(math.Sqrt(float64(bAvg)))
	a := uint8(math.Sqrt(float64(aAvg)))

	return color.RGBA{R: r, G: g, B: b, A: a}
}

func colorToPoint(clr color.RGBA) rtreego.Point {
	res := make([]float64, 4)

	r, g, b, a := rgba32ToRGBA8(clr.RGBA())

	rFloat := float64(r)
	gFloat := float64(g)
	bFloat := float64(b)
	aFloat := float64(a)

	res[0] = rFloat
	res[1] = gFloat
	res[2] = bFloat
	res[3] = aFloat

	return res
}