feat: wave-3 image + document surfaces — segmentation, colorize, face restore, ocr (ADR-0023)
- imagegen.Segmenter/SegmentationProvider: prompted mask via POST /upstream/<id>/v1/segment (file, prompt[, threshold], output=mask); white = prompted region, EditRequest.Mask polarity. - imagegen.Colorizer/ColorizeProvider: POST /upstream/<id>/v1/colorize. - imagegen.FaceRestorer/FaceRestoreProvider: POST /upstream/<id>/v1/restore_faces (upscale 1|2). - New ocr leaf package (Request/Page/Result, Recognize) + llamaswap OCRModel: POST /upstream/<id>/v1/ocr (file[, langs, max_pages]), tolerant per-page decode (join lines when page text absent), Raw escape hatch. - httptest contract tests per surface; ADR-0023; ADR index backfilled (0020-0022 rows were missing). Co-Authored-By: Claude Fable 5 <[email protected]>
This commit is contained in:
Claude Fable 5
@@ -0,0 +1,140 @@
|
||||
// ocr.go implements ocr.Provider against a Surya-style shim reached through
|
||||
// llama-swap's /upstream passthrough (ADR-0023):
|
||||
//
|
||||
// POST /upstream/<id>/v1/ocr multipart file[,langs,max_pages]
|
||||
//
|
||||
// The document may be an image or a PDF (the shim rasterizes PDFs itself).
|
||||
// The response is per-page JSON: {pages:[{number,text,lines,layout}]}. The
|
||||
// decode is tolerant — when a page carries no aggregate `text`, its line
|
||||
// texts are joined instead — and the full payload survives in Result.Raw for
|
||||
// callers that want bboxes/confidence/layout.
|
||||
package llamaswap
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"mime"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/ocr"
|
||||
)
|
||||
|
||||
// OCRModel implements ocr.Provider. The id selects which upstream llama-swap
|
||||
// loads.
|
||||
func (p *Provider) OCRModel(id string, opts ...ocr.ModelOption) (ocr.Model, error) {
|
||||
if err := p.requireBaseURL(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = ocr.ApplyModelOptions(opts)
|
||||
return &ocrModel{p: p, id: id}, nil
|
||||
}
|
||||
|
||||
type ocrModel struct {
|
||||
p *Provider
|
||||
id string
|
||||
}
|
||||
|
||||
// ocrResponse is the Surya shim's /v1/ocr shape (the subset this package
|
||||
// relies on; per-line bbox/confidence and layout stay in Raw).
|
||||
type ocrResponse struct {
|
||||
Pages []struct {
|
||||
Number int `json:"number"`
|
||||
Text string `json:"text"`
|
||||
Lines []struct {
|
||||
Text string `json:"text"`
|
||||
} `json:"lines"`
|
||||
} `json:"pages"`
|
||||
}
|
||||
|
||||
// Recognize implements ocr.Model.
|
||||
func (m *ocrModel) Recognize(ctx context.Context, req ocr.Request, opts ...ocr.Option) (*ocr.Result, error) {
|
||||
req = req.Apply(opts...)
|
||||
if len(req.Document) == 0 {
|
||||
return nil, fmt.Errorf("%w: ocr requires document bytes", llm.ErrUnsupported)
|
||||
}
|
||||
if req.MaxPages < 0 {
|
||||
return nil, fmt.Errorf("%w: ocr max pages must be >= 0, got %d", llm.ErrUnsupported, req.MaxPages)
|
||||
}
|
||||
path, err := upstreamPath(m.id, "/v1/ocr")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
maxPages := ""
|
||||
if req.MaxPages != 0 {
|
||||
maxPages = strconv.Itoa(req.MaxPages)
|
||||
}
|
||||
body, contentType, err := buildMultipart("build ocr form",
|
||||
filePart{field: "file", filename: documentFilename(req.Filename, req.MIME), data: req.Document},
|
||||
[]formField{
|
||||
{"langs", strings.Join(req.Languages, ","), false},
|
||||
{"max_pages", maxPages, false},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
raw, _, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxResponseBytes)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var out ocrResponse
|
||||
if err := json.Unmarshal(raw, &out); err != nil {
|
||||
return nil, fmt.Errorf("llama-swap: decode ocr response: %w", err)
|
||||
}
|
||||
if len(out.Pages) == 0 {
|
||||
// A blank page still comes back as a page with empty text; zero pages
|
||||
// is API drift or a soft error, never "the document was empty".
|
||||
return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
|
||||
Message: "ocr response contained no pages: " + truncateForError(raw)}
|
||||
}
|
||||
res := &ocr.Result{Raw: json.RawMessage(raw)}
|
||||
texts := make([]string, 0, len(out.Pages))
|
||||
for i, pg := range out.Pages {
|
||||
text := pg.Text
|
||||
if text == "" && len(pg.Lines) > 0 {
|
||||
lines := make([]string, 0, len(pg.Lines))
|
||||
for _, ln := range pg.Lines {
|
||||
lines = append(lines, ln.Text)
|
||||
}
|
||||
text = strings.Join(lines, "\n")
|
||||
}
|
||||
number := pg.Number
|
||||
if number == 0 {
|
||||
number = i + 1
|
||||
}
|
||||
res.Pages = append(res.Pages, ocr.Page{Number: number, Text: text})
|
||||
texts = append(texts, text)
|
||||
}
|
||||
res.Text = strings.Join(texts, "\n\n")
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// documentFilename picks the multipart filename hint for an OCR document: the
|
||||
// caller's (sanitized), else one derived from the MIME subtype
|
||||
// ("document.pdf"), else "document". MIME parameters are stripped before
|
||||
// matching, mirroring transcriptionFilename.
|
||||
func documentFilename(filename, mimeType string) string {
|
||||
if name := sanitizeFilename(filename); name != "" {
|
||||
return name
|
||||
}
|
||||
mt := strings.ToLower(strings.TrimSpace(mimeType))
|
||||
if parsed, _, err := mime.ParseMediaType(mt); err == nil {
|
||||
mt = parsed
|
||||
}
|
||||
switch mt {
|
||||
case "application/pdf":
|
||||
return "document.pdf"
|
||||
case "image/png":
|
||||
return "document.png"
|
||||
case "image/jpeg", "image/jpg":
|
||||
return "document.jpg"
|
||||
case "image/webp":
|
||||
return "document.webp"
|
||||
default:
|
||||
return "document"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
package llamaswap
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/ocr"
|
||||
)
|
||||
|
||||
func TestRecognize(t *testing.T) {
|
||||
var gotPath, gotLangs, gotMaxPages, gotFilename string
|
||||
var gotFile []byte
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotPath = r.URL.Path
|
||||
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
||||
t.Fatalf("parse multipart: %v", err)
|
||||
}
|
||||
gotLangs = r.FormValue("langs")
|
||||
gotMaxPages = r.FormValue("max_pages")
|
||||
f, hdr, err := r.FormFile("file")
|
||||
if err != nil {
|
||||
t.Fatalf("form file: %v", err)
|
||||
}
|
||||
defer f.Close()
|
||||
gotFile, _ = io.ReadAll(f)
|
||||
gotFilename = hdr.Filename
|
||||
_, _ = w.Write([]byte(`{"pages":[
|
||||
{"number":1,"text":"page one","lines":[{"text":"page","bbox":[0,0,1,1],"confidence":0.9},{"text":"one"}],"layout":{}},
|
||||
{"number":2,"text":"page two"}
|
||||
]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
om, err := p.OCRModel("ocr-surya")
|
||||
if err != nil {
|
||||
t.Fatalf("OCRModel: %v", err)
|
||||
}
|
||||
res, err := om.Recognize(context.Background(),
|
||||
ocr.Request{Document: []byte("%PDF"), MIME: "application/pdf"},
|
||||
ocr.WithLanguages("en", "de"), ocr.WithMaxPages(5))
|
||||
if err != nil {
|
||||
t.Fatalf("Recognize: %v", err)
|
||||
}
|
||||
if gotPath != "/upstream/ocr-surya/v1/ocr" {
|
||||
t.Errorf("path = %q", gotPath)
|
||||
}
|
||||
if gotLangs != "en,de" || gotMaxPages != "5" {
|
||||
t.Errorf("langs/max_pages = %q/%q", gotLangs, gotMaxPages)
|
||||
}
|
||||
if string(gotFile) != "%PDF" || gotFilename != "document.pdf" {
|
||||
t.Errorf("file = %q name = %q", gotFile, gotFilename)
|
||||
}
|
||||
if len(res.Pages) != 2 || res.Pages[0].Number != 1 || res.Pages[0].Text != "page one" ||
|
||||
res.Pages[1].Number != 2 || res.Pages[1].Text != "page two" {
|
||||
t.Errorf("pages = %+v", res.Pages)
|
||||
}
|
||||
if res.Text != "page one\n\npage two" {
|
||||
t.Errorf("text = %q", res.Text)
|
||||
}
|
||||
if res.Raw == nil {
|
||||
t.Error("Raw = nil, want raw payload")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecognizeJoinsLinesWhenPageTextAbsent(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write([]byte(`{"pages":[{"lines":[{"text":"first line"},{"text":"second line"}]}]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
om, _ := p.OCRModel("ocr-surya")
|
||||
res, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")})
|
||||
if err != nil {
|
||||
t.Fatalf("Recognize: %v", err)
|
||||
}
|
||||
if len(res.Pages) != 1 || res.Pages[0].Text != "first line\nsecond line" {
|
||||
t.Errorf("pages = %+v", res.Pages)
|
||||
}
|
||||
// Missing page number defaults to position.
|
||||
if res.Pages[0].Number != 1 {
|
||||
t.Errorf("number = %d, want 1", res.Pages[0].Number)
|
||||
}
|
||||
if res.Text != "first line\nsecond line" {
|
||||
t.Errorf("text = %q", res.Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecognizeOmitsUnsetFields(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
||||
t.Fatalf("parse multipart: %v", err)
|
||||
}
|
||||
for _, k := range []string{"langs", "max_pages"} {
|
||||
if v, ok := r.MultipartForm.Value[k]; ok {
|
||||
t.Errorf("unset request sent %q = %v, want omitted", k, v)
|
||||
}
|
||||
}
|
||||
if _, hdr, err := r.FormFile("file"); err == nil {
|
||||
if hdr.Filename != "document" {
|
||||
t.Errorf("filename = %q, want document fallback", hdr.Filename)
|
||||
}
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"pages":[{"number":1,"text":"x"}]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
om, _ := p.OCRModel("ocr-surya")
|
||||
if _, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")}); err != nil {
|
||||
t.Fatalf("Recognize: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecognizeRejectsBadArgs(t *testing.T) {
|
||||
p := New(WithBaseURL("http://unused"))
|
||||
om, _ := p.OCRModel("ocr-surya")
|
||||
if _, err := om.Recognize(context.Background(), ocr.Request{}); !errors.Is(err, llm.ErrUnsupported) {
|
||||
t.Errorf("no document: err = %v, want ErrUnsupported", err)
|
||||
}
|
||||
if _, err := om.Recognize(context.Background(),
|
||||
ocr.Request{Document: []byte{1}, MaxPages: -1}); !errors.Is(err, llm.ErrUnsupported) {
|
||||
t.Errorf("negative max pages: err = %v, want ErrUnsupported", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecognizeRejectsZeroPages(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write([]byte(`{"pages":[]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
om, _ := p.OCRModel("ocr-surya")
|
||||
_, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")})
|
||||
var apiErr *llm.APIError
|
||||
if !errors.As(err, &apiErr) {
|
||||
t.Fatalf("err = %v, want APIError for zero pages", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecognizeRejectsNonJSONResponse(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/html")
|
||||
_, _ = w.Write([]byte("<html>proxy error page</html>"))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
om, _ := p.OCRModel("ocr-surya")
|
||||
if _, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")}); err == nil ||
|
||||
!strings.Contains(err.Error(), "decode ocr response") {
|
||||
t.Fatalf("err = %v, want decode error for non-JSON body", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecognizeSurfacesAPIError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
_, _ = w.Write([]byte(`{"error":{"message":"unsupported file type"}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
om, _ := p.OCRModel("ocr-surya")
|
||||
_, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("bad")})
|
||||
var apiErr *llm.APIError
|
||||
if !errors.As(err, &apiErr) {
|
||||
t.Fatalf("err = %T %v, want *llm.APIError", err, err)
|
||||
}
|
||||
if apiErr.Status != http.StatusBadRequest || apiErr.Message != "unsupported file type" || apiErr.Model != "ocr-surya" {
|
||||
t.Errorf("apiErr = %+v", apiErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentFilename(t *testing.T) {
|
||||
cases := []struct {
|
||||
filename, mime, want string
|
||||
}{
|
||||
{"scan.pdf", "", "scan.pdf"},
|
||||
{"evil\r\nname.pdf", "", "evilname.pdf"},
|
||||
{"", "application/pdf", "document.pdf"},
|
||||
{"", "image/png", "document.png"},
|
||||
{"", "image/jpeg", "document.jpg"},
|
||||
{"", "image/webp; charset=binary", "document.webp"},
|
||||
{"", "", "document"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
if got := documentFilename(tc.filename, tc.mime); got != tc.want {
|
||||
t.Errorf("documentFilename(%q, %q) = %q, want %q", tc.filename, tc.mime, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
// restore.go implements the imagegen.ColorizeProvider and
|
||||
// imagegen.FaceRestoreProvider surfaces against the mediautils shim reached
|
||||
// through llama-swap's /upstream passthrough (ADR-0023):
|
||||
//
|
||||
// colorize POST /upstream/<id>/v1/colorize (DDColor)
|
||||
// restore_faces POST /upstream/<id>/v1/restore_faces (GFPGAN)
|
||||
//
|
||||
// Both are one-file multipart in, one PNG out, mirroring mediautil.go.
|
||||
package llamaswap
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/imagegen"
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
|
||||
)
|
||||
|
||||
// --- colorize ---
|
||||
|
||||
// ColorizeModel implements imagegen.ColorizeProvider against the mediautils
|
||||
// shim's POST /v1/colorize. The id selects which upstream llama-swap loads.
|
||||
func (p *Provider) ColorizeModel(id string, opts ...imagegen.ColorizeModelOption) (imagegen.Colorizer, error) {
|
||||
if err := p.requireBaseURL(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = imagegen.ApplyColorizeModelOptions(opts)
|
||||
return &colorizeModel{p: p, id: id}, nil
|
||||
}
|
||||
|
||||
type colorizeModel struct {
|
||||
p *Provider
|
||||
id string
|
||||
}
|
||||
|
||||
// Colorize implements imagegen.Colorizer.
|
||||
func (m *colorizeModel) Colorize(ctx context.Context, req imagegen.ColorizeRequest, opts ...imagegen.ColorizeOption) (*imagegen.Result, error) {
|
||||
req = req.Apply(opts...)
|
||||
if len(req.Image.Data) == 0 {
|
||||
return nil, fmt.Errorf("%w: colorization requires an image", llm.ErrUnsupported)
|
||||
}
|
||||
path, err := upstreamPath(m.id, "/v1/colorize")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body, contentType, err := buildMultipart("build colorize form",
|
||||
filePart{field: "file", filename: "image.png", data: req.Image.Data},
|
||||
nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxImageResponseBytes)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return singleImageResult(m.p.name, m.id, "colorize", raw, respType)
|
||||
}
|
||||
|
||||
// --- face restoration ---
|
||||
|
||||
// FaceRestoreModel implements imagegen.FaceRestoreProvider against the
|
||||
// mediautils shim's POST /v1/restore_faces.
|
||||
func (p *Provider) FaceRestoreModel(id string, opts ...imagegen.FaceRestoreModelOption) (imagegen.FaceRestorer, error) {
|
||||
if err := p.requireBaseURL(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = imagegen.ApplyFaceRestoreModelOptions(opts)
|
||||
return &faceRestoreModel{p: p, id: id}, nil
|
||||
}
|
||||
|
||||
type faceRestoreModel struct {
|
||||
p *Provider
|
||||
id string
|
||||
}
|
||||
|
||||
// RestoreFaces implements imagegen.FaceRestorer.
|
||||
func (m *faceRestoreModel) RestoreFaces(ctx context.Context, req imagegen.FaceRestoreRequest, opts ...imagegen.FaceRestoreOption) (*imagegen.Result, error) {
|
||||
req = req.Apply(opts...)
|
||||
if len(req.Image.Data) == 0 {
|
||||
return nil, fmt.Errorf("%w: face restoration requires an image", llm.ErrUnsupported)
|
||||
}
|
||||
if req.Upscale != 0 && req.Upscale != 1 && req.Upscale != 2 {
|
||||
return nil, fmt.Errorf("%w: face-restore upscale must be 1 or 2, got %d", llm.ErrUnsupported, req.Upscale)
|
||||
}
|
||||
path, err := upstreamPath(m.id, "/v1/restore_faces")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
upscale := ""
|
||||
if req.Upscale != 0 {
|
||||
upscale = strconv.Itoa(req.Upscale)
|
||||
}
|
||||
body, contentType, err := buildMultipart("build restore-faces form",
|
||||
filePart{field: "file", filename: "image.png", data: req.Image.Data},
|
||||
[]formField{{"upscale", upscale, false}})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxImageResponseBytes)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return singleImageResult(m.p.name, m.id, "face restoration", raw, respType)
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
// segment.go implements imagegen.SegmentationProvider against a
|
||||
// GroundingDINO+SAM shim (segment-langsam) reached through llama-swap's
|
||||
// /upstream passthrough (ADR-0023):
|
||||
//
|
||||
// POST /upstream/<id>/v1/segment multipart file,prompt[,threshold],output=mask
|
||||
//
|
||||
// The response is a single grayscale mask PNG where WHITE marks the prompted
|
||||
// region — directly usable as imagegen.EditRequest.Mask (white = repaint).
|
||||
// The shim also offers output=cutout|boxes; this client always requests the
|
||||
// mask, because a cutout is derivable client-side from mask+original with no
|
||||
// second GPU call.
|
||||
package llamaswap
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/imagegen"
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
|
||||
)
|
||||
|
||||
// SegmentationModel implements imagegen.SegmentationProvider. The id selects
|
||||
// which upstream llama-swap loads.
|
||||
func (p *Provider) SegmentationModel(id string, opts ...imagegen.SegmentationModelOption) (imagegen.Segmenter, error) {
|
||||
if err := p.requireBaseURL(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = imagegen.ApplySegmentationModelOptions(opts)
|
||||
return &segmentationModel{p: p, id: id}, nil
|
||||
}
|
||||
|
||||
type segmentationModel struct {
|
||||
p *Provider
|
||||
id string
|
||||
}
|
||||
|
||||
// Segment implements imagegen.Segmenter.
|
||||
func (m *segmentationModel) Segment(ctx context.Context, req imagegen.SegmentationRequest, opts ...imagegen.SegmentationOption) (*imagegen.Result, error) {
|
||||
req = req.Apply(opts...)
|
||||
if len(req.Image.Data) == 0 {
|
||||
return nil, fmt.Errorf("%w: segmentation requires an image", llm.ErrUnsupported)
|
||||
}
|
||||
if req.Prompt == "" {
|
||||
return nil, fmt.Errorf("%w: segmentation requires a prompt", llm.ErrUnsupported)
|
||||
}
|
||||
if req.Threshold < 0 || req.Threshold > 1 {
|
||||
return nil, fmt.Errorf("%w: segmentation threshold must be in [0,1], got %g", llm.ErrUnsupported, req.Threshold)
|
||||
}
|
||||
path, err := upstreamPath(m.id, "/v1/segment")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
threshold := ""
|
||||
if req.Threshold != 0 {
|
||||
threshold = strconv.FormatFloat(req.Threshold, 'g', -1, 64)
|
||||
}
|
||||
body, contentType, err := buildMultipart("build segment form",
|
||||
filePart{field: "file", filename: "image.png", data: req.Image.Data},
|
||||
[]formField{
|
||||
{"prompt", req.Prompt, true},
|
||||
{"threshold", threshold, false},
|
||||
// Always the mask: white = prompted region, EditRequest.Mask
|
||||
// polarity. Cutouts are derived client-side.
|
||||
{"output", "mask", true},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxImageResponseBytes)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return singleImageResult(m.p.name, m.id, "segmentation", raw, respType)
|
||||
}
|
||||
@@ -0,0 +1,270 @@
|
||||
package llamaswap
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/imagegen"
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
|
||||
)
|
||||
|
||||
func TestSegment(t *testing.T) {
|
||||
png := pngFixture(t)
|
||||
var gotPath, gotPrompt, gotThreshold, gotOutput, gotFilename string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotPath = r.URL.Path
|
||||
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
||||
t.Errorf("parse multipart: %v", err)
|
||||
}
|
||||
gotPrompt = r.FormValue("prompt")
|
||||
gotThreshold = r.FormValue("threshold")
|
||||
gotOutput = r.FormValue("output")
|
||||
if f, hdr, err := r.FormFile("file"); err == nil {
|
||||
gotFilename = hdr.Filename
|
||||
f.Close()
|
||||
} else {
|
||||
t.Errorf("file part: %v", err)
|
||||
}
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(png)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
sg, err := p.SegmentationModel("segment-langsam")
|
||||
if err != nil {
|
||||
t.Fatalf("SegmentationModel: %v", err)
|
||||
}
|
||||
res, err := sg.Segment(context.Background(),
|
||||
imagegen.SegmentationRequest{Image: imagegen.Image{MIME: "image/png", Data: png}, Prompt: "the red car"},
|
||||
imagegen.WithSegmentationThreshold(0.35))
|
||||
if err != nil {
|
||||
t.Fatalf("Segment: %v", err)
|
||||
}
|
||||
if gotPath != "/upstream/segment-langsam/v1/segment" {
|
||||
t.Errorf("path = %q", gotPath)
|
||||
}
|
||||
if gotPrompt != "the red car" || gotThreshold != "0.35" || gotOutput != "mask" || gotFilename != "image.png" {
|
||||
t.Errorf("prompt/threshold/output/filename = %q/%q/%q/%q", gotPrompt, gotThreshold, gotOutput, gotFilename)
|
||||
}
|
||||
if len(res.Images) != 1 || res.Images[0].MIME != "image/png" {
|
||||
t.Fatalf("images = %+v", res.Images)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSegmentOmitsDefaultThreshold(t *testing.T) {
|
||||
png := pngFixture(t)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
||||
t.Errorf("parse multipart: %v", err)
|
||||
}
|
||||
if _, ok := r.MultipartForm.Value["threshold"]; ok {
|
||||
t.Error("threshold field sent for default; want omitted")
|
||||
}
|
||||
if got := r.FormValue("output"); got != "mask" {
|
||||
t.Errorf("output = %q, want mask (always sent)", got)
|
||||
}
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(png)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
sg, _ := p.SegmentationModel("segment-langsam")
|
||||
if _, err := sg.Segment(context.Background(),
|
||||
imagegen.SegmentationRequest{Image: imagegen.Image{Data: png}, Prompt: "dog"}); err != nil {
|
||||
t.Fatalf("Segment: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSegmentRejectsBadArgs(t *testing.T) {
|
||||
p := New(WithBaseURL("http://unused"))
|
||||
sg, _ := p.SegmentationModel("segment-langsam")
|
||||
if _, err := sg.Segment(context.Background(),
|
||||
imagegen.SegmentationRequest{Prompt: "dog"}); !errors.Is(err, llm.ErrUnsupported) {
|
||||
t.Errorf("no image: err = %v, want ErrUnsupported", err)
|
||||
}
|
||||
if _, err := sg.Segment(context.Background(),
|
||||
imagegen.SegmentationRequest{Image: imagegen.Image{Data: []byte{1}}}); !errors.Is(err, llm.ErrUnsupported) {
|
||||
t.Errorf("no prompt: err = %v, want ErrUnsupported", err)
|
||||
}
|
||||
if _, err := sg.Segment(context.Background(),
|
||||
imagegen.SegmentationRequest{Image: imagegen.Image{Data: []byte{1}}, Prompt: "dog", Threshold: 1.5}); !errors.Is(err, llm.ErrUnsupported) {
|
||||
t.Errorf("threshold 1.5: err = %v, want ErrUnsupported", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSegmentRejectsNonImageResponse(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/html")
|
||||
_, _ = w.Write([]byte("<html>proxy error page</html>"))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
sg, _ := p.SegmentationModel("segment-langsam")
|
||||
_, err := sg.Segment(context.Background(),
|
||||
imagegen.SegmentationRequest{Image: imagegen.Image{Data: pngFixture(t)}, Prompt: "dog"})
|
||||
var apiErr *llm.APIError
|
||||
if !errors.As(err, &apiErr) {
|
||||
t.Fatalf("err = %v, want APIError for non-image body", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestColorize(t *testing.T) {
|
||||
png := pngFixture(t)
|
||||
var gotPath, gotFilename string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotPath = r.URL.Path
|
||||
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
||||
t.Errorf("parse multipart: %v", err)
|
||||
}
|
||||
if f, hdr, err := r.FormFile("file"); err == nil {
|
||||
gotFilename = hdr.Filename
|
||||
f.Close()
|
||||
} else {
|
||||
t.Errorf("file part: %v", err)
|
||||
}
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(png)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
cl, err := p.ColorizeModel("mediautils")
|
||||
if err != nil {
|
||||
t.Fatalf("ColorizeModel: %v", err)
|
||||
}
|
||||
res, err := cl.Colorize(context.Background(),
|
||||
imagegen.ColorizeRequest{Image: imagegen.Image{Data: png}})
|
||||
if err != nil {
|
||||
t.Fatalf("Colorize: %v", err)
|
||||
}
|
||||
if gotPath != "/upstream/mediautils/v1/colorize" {
|
||||
t.Errorf("path = %q", gotPath)
|
||||
}
|
||||
if gotFilename != "image.png" {
|
||||
t.Errorf("filename = %q", gotFilename)
|
||||
}
|
||||
if len(res.Images) != 1 || res.Images[0].MIME != "image/png" {
|
||||
t.Fatalf("images = %+v", res.Images)
|
||||
}
|
||||
}
|
||||
|
||||
func TestColorizeRejectsEmptyImage(t *testing.T) {
|
||||
p := New(WithBaseURL("http://unused"))
|
||||
cl, _ := p.ColorizeModel("mediautils")
|
||||
if _, err := cl.Colorize(context.Background(), imagegen.ColorizeRequest{}); !errors.Is(err, llm.ErrUnsupported) {
|
||||
t.Errorf("err = %v, want ErrUnsupported", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestColorizeSurfacesAPIError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
_, _ = w.Write([]byte(`{"error":{"message":"model is loading"}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
cl, _ := p.ColorizeModel("mediautils")
|
||||
_, err := cl.Colorize(context.Background(),
|
||||
imagegen.ColorizeRequest{Image: imagegen.Image{Data: pngFixture(t)}})
|
||||
var apiErr *llm.APIError
|
||||
if !errors.As(err, &apiErr) {
|
||||
t.Fatalf("err = %T %v, want *llm.APIError", err, err)
|
||||
}
|
||||
if apiErr.Status != http.StatusServiceUnavailable || apiErr.Message != "model is loading" || apiErr.Model != "mediautils" {
|
||||
t.Errorf("apiErr = %+v", apiErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreFaces(t *testing.T) {
|
||||
png := pngFixture(t)
|
||||
var gotPath, gotUpscale string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotPath = r.URL.Path
|
||||
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
||||
t.Errorf("parse multipart: %v", err)
|
||||
}
|
||||
gotUpscale = r.FormValue("upscale")
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(png)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
fr, err := p.FaceRestoreModel("mediautils")
|
||||
if err != nil {
|
||||
t.Fatalf("FaceRestoreModel: %v", err)
|
||||
}
|
||||
res, err := fr.RestoreFaces(context.Background(),
|
||||
imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: png}},
|
||||
imagegen.WithFaceRestoreUpscale(2))
|
||||
if err != nil {
|
||||
t.Fatalf("RestoreFaces: %v", err)
|
||||
}
|
||||
if gotPath != "/upstream/mediautils/v1/restore_faces" {
|
||||
t.Errorf("path = %q", gotPath)
|
||||
}
|
||||
if gotUpscale != "2" {
|
||||
t.Errorf("upscale = %q", gotUpscale)
|
||||
}
|
||||
if len(res.Images) != 1 {
|
||||
t.Fatalf("images = %+v", res.Images)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreFacesOmitsDefaultUpscale(t *testing.T) {
|
||||
png := pngFixture(t)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
||||
t.Errorf("parse multipart: %v", err)
|
||||
}
|
||||
if _, ok := r.MultipartForm.Value["upscale"]; ok {
|
||||
t.Error("upscale field sent for default; want omitted")
|
||||
}
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(png)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
fr, _ := p.FaceRestoreModel("mediautils")
|
||||
if _, err := fr.RestoreFaces(context.Background(),
|
||||
imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: png}}); err != nil {
|
||||
t.Fatalf("RestoreFaces: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreFacesRejectsBadArgs(t *testing.T) {
|
||||
p := New(WithBaseURL("http://unused"))
|
||||
fr, _ := p.FaceRestoreModel("mediautils")
|
||||
if _, err := fr.RestoreFaces(context.Background(), imagegen.FaceRestoreRequest{}); !errors.Is(err, llm.ErrUnsupported) {
|
||||
t.Errorf("no image: err = %v, want ErrUnsupported", err)
|
||||
}
|
||||
if _, err := fr.RestoreFaces(context.Background(),
|
||||
imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: []byte{1}}, Upscale: 3}); !errors.Is(err, llm.ErrUnsupported) {
|
||||
t.Errorf("upscale 3: err = %v, want ErrUnsupported", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreFacesRejectsNonImageResponse(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header()["Content-Type"] = nil // NO Content-Type at all
|
||||
_, _ = w.Write([]byte("502 bad gateway"))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
fr, _ := p.FaceRestoreModel("mediautils")
|
||||
_, err := fr.RestoreFaces(context.Background(),
|
||||
imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: pngFixture(t)}})
|
||||
var apiErr *llm.APIError
|
||||
if !errors.As(err, &apiErr) {
|
||||
t.Fatalf("err = %v, want APIError for headerless non-image body", err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user