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:
@@ -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"
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user