Maven/internal/media/image.go

package media

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"image"
	"image/gif"
	"image/jpeg"
	"image/png"
	"strings"
)

// DefaultMaxDim — the longest edge an image is scaled down to before it goes to
// a vision model. 896 is the tile size the current crop of small
// vision-language models (Qwen2.5-VL, SmolVLM, moondream) work in; sending a
// 12-megapixel phone photo instead just costs the box minutes of prefill for
// tiles that get pooled away anyway.
const DefaultMaxDim = 896

// JPEGQuality for the re-encode. 85 is the usual "no visible artefacts" point,
// and the re-encode exists to shrink the payload, not to archive it — the
// original bytes stay in the blob store untouched.
const JPEGQuality = 85

// DefaultMaxPixels — the largest source image this build will decode, counted
// in pixels rather than in compressed bytes. A byte cap is not a memory bound
// for an image: a 20000x20000 PNG of flat colour compresses to a few hundred
// kilobytes and decodes to 400 million pixels, which is 1.6 GB of heap in the
// process that owns the database and the socket. 40 megapixels is well past any
// phone camera and two orders of magnitude short of an OOM.
const DefaultMaxPixels = 40 << 20

// ErrTooManyPixels — the image header declares more pixels than this build
// will decode. Separate from ErrUnsupportedImage because the format is fine and
// the size is not, and the log line should say which.
var ErrTooManyPixels = errors.New("media: image has too many pixels")

// ErrUnsupportedImage — the bytes are not an image format this build can
// decode. Notably webp: the stdlib has no webp decoder and this repo takes no
// new dependencies, so a webp arriving from Telegram is refused here with a
// clear error rather than handed to a model as garbage.
var ErrUnsupportedImage = errors.New("media: unsupported image format")

// SniffImage identifies image bytes by magic number and returns the mime. It
// exists because a caller-declared content type is a claim, and the store's file
// extension (and the vision provider's data URI) should follow the bytes.
//
// Returns ErrUnsupportedImage for anything unrecognised, including webp — which
// is recognised well enough to name in the error, so the log says "webp is not
// supported" instead of "not an image".
func SniffImage(data []byte) (string, error) {
	switch {
	case len(data) >= 3 && data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF:
		return "image/jpeg", nil
	case len(data) >= 8 && string(data[:8]) == "\x89PNG\r\n\x1a\n":
		return "image/png", nil
	case len(data) >= 6 && (string(data[:6]) == "GIF87a" || string(data[:6]) == "GIF89a"):
		return "image/gif", nil
	case len(data) >= 12 && string(data[:4]) == "RIFF" && string(data[8:12]) == "WEBP":
		return "", fmt.Errorf("%w: webp (no decoder in this build)", ErrUnsupportedImage)
	}
	return "", ErrUnsupportedImage
}

// Image — an image prepared for a vision model: JPEG bytes, downscaled, with
// the dimensions it ended up at. It is deliberately a separate type from Blob:
// a Blob is what he sent, an Image is what the model sees, and the two are not
// the same bytes.
type Image struct {
	JPEG   []byte
	Width  int
	Height int
	// Source names where the original came from ("telegram", "web:upload"),
	// carried through only so a log line can say what was looked at.
	Source string
}

// DataURI renders the image as a `data:image/jpeg;base64,...` URI, which is how
// every OpenAI-compatible multimodal endpoint takes an image. The string is
// large (roughly 4/3 of the JPEG); nothing caches it.
func (im Image) DataURI() string {
	return "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(im.JPEG)
}

// PrepareImage decodes data, scales it so its longest edge is at most maxDim
// (never up — a small image is left alone), and re-encodes it as JPEG.
// maxDim ≤ 0 ⇒ DefaultMaxDim.
//
// An image with an alpha channel is composited onto white rather than having
// alpha dropped to black, because the common case is a screenshot or a
// transparent-background diagram, and text on black-on-black is unreadable to
// the model for no reason.
func PrepareImage(data []byte, source string, maxDim int) (Image, error) {
	if len(data) == 0 {
		return Image{}, ErrEmpty
	}
	if maxDim <= 0 {
		maxDim = DefaultMaxDim
	}
	mime, err := SniffImage(data)
	if err != nil {
		return Image{}, err
	}
	// The header is read before the pixels. Deciding after the decode is not a
	// decision: by then the whole bitmap is already in the heap.
	cfg, err := decodeConfig(data, mime)
	if err != nil {
		return Image{}, fmt.Errorf("media: read %s header: %w", mime, err)
	}
	if px := int64(cfg.Width) * int64(cfg.Height); px > DefaultMaxPixels {
		return Image{}, fmt.Errorf("%w: %dx%d is %d, cap is %d",
			ErrTooManyPixels, cfg.Width, cfg.Height, px, int64(DefaultMaxPixels))
	}
	src, err := decode(data, mime)
	if err != nil {
		return Image{}, fmt.Errorf("media: decode %s: %w", mime, err)
	}

	dst := flattenAndScale(src, maxDim)
	var buf bytes.Buffer
	if err := jpeg.Encode(&buf, dst, &jpeg.Options{Quality: JPEGQuality}); err != nil {
		return Image{}, fmt.Errorf("media: encode jpeg: %w", err)
	}
	b := dst.Bounds()
	return Image{JPEG: buf.Bytes(), Width: b.Dx(), Height: b.Dy(), Source: source}, nil
}

func decodeConfig(data []byte, mime string) (image.Config, error) {
	r := bytes.NewReader(data)
	switch strings.ToLower(mime) {
	case "image/jpeg":
		return jpeg.DecodeConfig(r)
	case "image/png":
		return png.DecodeConfig(r)
	case "image/gif":
		return gif.DecodeConfig(r)
	}
	return image.Config{}, ErrUnsupportedImage
}

func decode(data []byte, mime string) (image.Image, error) {
	r := bytes.NewReader(data)
	switch strings.ToLower(mime) {
	case "image/jpeg":
		return jpeg.Decode(r)
	case "image/png":
		return png.Decode(r)
	case "image/gif":
		return gif.Decode(r)
	}
	return nil, ErrUnsupportedImage
}

// flattenAndScale composites onto white and box-scales down to maxDim. The
// scaler is a plain area average over the source pixels mapping to each
// destination pixel — nearest-neighbour would alias small text into noise,
// which defeats the point of reading a screenshot, and an area average is a
// dozen lines against pulling in golang.org/x/image on an offline box.
//
// It reads the source through At and allocates only the destination. Flattening
// into a full-size RGBA first doubled the peak: a 40-megapixel photo already
// costs 160 MB decoded, and the intermediate made it 320 MB before MaxDim had
// any chance to help.
func flattenAndScale(src image.Image, maxDim int) *image.RGBA {
	sb := src.Bounds()
	sw, sh := sb.Dx(), sb.Dy()
	dw, dh := fit(sw, sh, maxDim)

	dst := image.NewRGBA(image.Rect(0, 0, dw, dh))
	for y := 0; y < dh; y++ {
		y0, y1 := y*sh/dh, (y+1)*sh/dh
		if y1 <= y0 {
			y1 = y0 + 1
		}
		for x := 0; x < dw; x++ {
			x0, x1 := x*sw/dw, (x+1)*sw/dw
			if x1 <= x0 {
				x1 = x0 + 1
			}
			var r, g, b, n uint64
			for sy := y0; sy < y1; sy++ {
				for sx := x0; sx < x1; sx++ {
					// At returns premultiplied 16-bit. Compositing over white
					// is then c + (1-alpha), which is the same answer the
					// draw.Over pass used to give, one pixel at a time.
					cr, cg, cb, ca := src.At(sb.Min.X+sx, sb.Min.Y+sy).RGBA()
					inv := uint64(0xFFFF - ca)
					r += uint64(cr) + inv
					g += uint64(cg) + inv
					b += uint64(cb) + inv
					n++
				}
			}
			o := dst.PixOffset(x, y)
			dst.Pix[o] = uint8(r / n >> 8)
			dst.Pix[o+1] = uint8(g / n >> 8)
			dst.Pix[o+2] = uint8(b / n >> 8)
			dst.Pix[o+3] = 0xFF
		}
	}
	return dst
}

// fit returns the largest w×h with the same aspect ratio whose longest edge is
// at most maxDim, never enlarging. Both edges are clamped to at least 1 so a
// 2000×1 strip does not scale to zero height.
func fit(w, h, maxDim int) (int, int) {
	if w <= maxDim && h <= maxDim {
		return w, h
	}
	if w >= h {
		nh := h * maxDim / w
		if nh < 1 {
			nh = 1
		}
		return maxDim, nh
	}
	nw := w * maxDim / h
	if nw < 1 {
		nw = 1
	}
	return nw, maxDim
}