Draw mouse in µniverse recording

main.go

// Command add_mouse adds a small square to every
// observation in a recording to indicate where the
// mouse was in that frame.
package main

import (
	"image"
	"image/color"
	"os"

	"github.com/unixpickle/essentials"
	"github.com/unixpickle/muniverse"
	"github.com/unixpickle/muniverse/chrome"
)

const MouseSize = 5

func main() {
	if len(os.Args) != 3 {
		essentials.Die("Usage: add_mouse <recording_in> <recording_out>")
	}
	recording, err := muniverse.OpenRecording(os.Args[1])
	if err != nil {
		essentials.Die(err)
	}
	output, err := muniverse.CreateRecording(os.Args[2])
	if err != nil {
		essentials.Die(err)
	}
	mouseX := -10
	mouseY := -10
	for i := 0; i < recording.NumObs(); i++ {
		obs, err := recording.ReadObs(i)
		if err != nil {
			essentials.Die(err)
		}
		newObs := drawMouse(obs, mouseX, mouseY)
		if err := output.WriteObs(newObs); err != nil {
			essentials.Die(err)
		}
		if i < recording.NumSteps() {
			step, err := recording.ReadStep(i)
			if err != nil {
				essentials.Die(err)
			}
			if err := output.WriteStep(step); err != nil {
				essentials.Die(err)
			}
			for _, evt := range step.Events {
				if mouse, ok := evt.(*chrome.MouseEvent); ok {
					mouseX = mouse.X
					mouseY = mouse.Y
				}
			}
		}
	}
}

func drawMouse(obs muniverse.Obs, x, y int) muniverse.Obs {
	img, err := obs.Image()
	if err != nil {
		essentials.Die(err)
	}
	setter := img.(ColorSetter)
	for subX := x - MouseSize; subX <= x+MouseSize; subX++ {
		for subY := y - MouseSize; subY <= y+MouseSize; subY++ {
			setter.Set(subX, subY, color.Gray{Y: 0xff})
		}
	}
	return Obs{img: img}
}

type ColorSetter interface {
	image.Image
	Set(x, y int, c color.Color)
}

type Obs struct {
	img image.Image
}

func (o Obs) Image() (image.Image, error) {
	return o.img, nil
}