ollama/model/qwen2vl/imageproc.go
Patrick Devine 8c9fb8eb73
imageproc mllama refactor (#7537)
Refactor mllama image processing code, and add pixtral and qwen2vl
2024-12-14 19:50:15 -08:00

75 lines
2.0 KiB
Go

package qwen2vl
import (
"fmt"
"image"
_ "image/jpeg"
_ "image/png"
"io"
"math"
"github.com/ollama/ollama/model/imageproc"
)
const (
DefaultFactor = 28
DefaultMinPixels = 56 * 56
DefaultMaxPixels = 14 * 14 * 4 * 1280
)
// smartResize calculates the size of the image to resize to based on the
// factor, minPixels, and maxPixels.
func smartResize(size image.Point, factor, minPixels, maxPixels int) image.Point {
// 1. Both dimensions of size are divisible by factor
// 2. The area of the image is between minPixels and maxPixels
// 3. The aspect ratio of the image is as close to 1:1 as possible
if size.Y < factor || size.X < factor {
panic("image is too small to resize")
} else if max(size.X, size.Y)/min(size.X, size.Y) > 200 {
panic("aspect ratio must be less than 200:1")
}
f := float64(factor)
width := float64(size.X)
height := float64(size.Y)
xBar := math.Round(width/f) * f
yBar := math.Round(height/f) * f
if xBar*yBar > float64(maxPixels) {
beta := math.Sqrt(height * width / float64(maxPixels))
xBar = math.Floor(width/beta/f) * f
yBar = math.Floor(height/beta/f) * f
} else if xBar*yBar < float64(minPixels) {
beta := math.Sqrt(float64(minPixels) / (height * width))
xBar = math.Ceil(width*beta/f) * f
yBar = math.Ceil(height*beta/f) * f
}
return image.Point{int(xBar), int(yBar)}
}
func resizeImage(img image.Image, format string, size image.Point) image.Image {
if format == "png" {
img = imageproc.Composite(img)
}
return imageproc.Resize(img, size, imageproc.ResizeBilinear)
}
func Preprocess(imageData io.Reader) ([]float32, map[string]any, error) {
img, format, err := image.Decode(imageData)
if err != nil {
return nil, nil, fmt.Errorf("failed to decode image: %w", err)
}
size := smartResize(img.Bounds().Max, DefaultFactor, DefaultMinPixels, DefaultMaxPixels)
img = resizeImage(img, format, size)
data := imageproc.Normalize(img, imageproc.ClipDefaultMean, imageproc.ClipDefaultSTD, true, true)
opts := map[string]any{}
return data, opts, nil
}