matmanna.dev / pic-to-patch
fork of gh:maxwoffard/pic-to-patch
0
Fork 0

Configure Feed

Select the types of activity you want to include in your feed.

Switch web to using go / cleanup worker build

author
Max Wofford date (May 11, 2026, 6:56 PM -0400) commit c55c1b22 parent 17778de8
+1457 -207
-32
Dockerfile
··· 1 - FROM python:3.11-slim 2 - 3 - RUN apt-get update && apt-get install -y --no-install-recommends \ 4 - wget xz-utils ca-certificates \ 5 - libgtk-3-0t64 libdrm2 libgbm1 libxcb-dri3-0 \ 6 - inkscape \ 7 - && rm -rf /var/lib/apt/lists/* 8 - 9 - # install inkstitch (detect arch) 10 - ARG INKSTITCH_VERSION=3.2.2 11 - RUN ARCH=$(dpkg --print-architecture) && \ 12 - if [ "$ARCH" = "arm64" ]; then INKARCH="aarch64"; else INKARCH="x86_64"; fi && \ 13 - wget -q "https://github.com/inkstitch/inkstitch/releases/download/v${INKSTITCH_VERSION}/inkstitch-${INKSTITCH_VERSION}-linux-${INKARCH}.tar.xz" \ 14 - -O /tmp/inkstitch.tar.xz && \ 15 - mkdir -p /opt/inkstitch && \ 16 - tar xJf /tmp/inkstitch.tar.xz -C /opt/inkstitch && \ 17 - rm /tmp/inkstitch.tar.xz 18 - ENV INKSTITCH_BIN=/opt/inkstitch/inkstitch/bin/inkstitch 19 - 20 - WORKDIR /app 21 - COPY requirements.txt . 22 - RUN pip install --no-cache-dir -r requirements.txt vtracer 23 - 24 - COPY pipeline/ pipeline/ 25 - COPY server.py worker.py ./ 26 - 27 - ENV RESULTS_DIR=/data/results 28 - ENV VTRACER_BIN=vtracer 29 - 30 - ENV MODE=api 31 - EXPOSE 8000 32 - CMD sh -c 'if [ "$MODE" = "worker" ]; then rq worker patches --url "$REDIS_URL"; else uvicorn server:app --host 0.0.0.0 --port 8000; fi'
+9 -13
Dockerfile.web
··· 1 - FROM python:3.11-slim 2 - 3 - WORKDIR /app 4 - 5 - RUN pip install --no-cache-dir \ 6 - fastapi==0.115.0 \ 7 - uvicorn==0.30.0 \ 8 - python-multipart==0.0.9 \ 9 - redis==5.0.0 \ 10 - rq==1.16.0 1 + FROM golang:1.24-alpine AS build 11 2 12 - COPY server.py . 3 + WORKDIR /src 4 + COPY web/go.mod web/go.sum ./ 5 + RUN go mod download 6 + COPY web/main.go . 7 + RUN CGO_ENABLED=0 go build -o /server . 13 8 14 - ENV RESULTS_DIR=/data/results 9 + FROM alpine:3.20 10 + COPY --from=build /server /server 15 11 EXPOSE 8000 16 - CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8000"] 12 + CMD ["/server"]
+24 -9
Dockerfile.worker
··· 1 - FROM python:3.11-slim 1 + # Stage 1: download and strip inkstitch 2 + FROM python:3.11-slim AS builder 2 3 3 4 RUN apt-get update && apt-get install -y --no-install-recommends \ 4 5 wget xz-utils ca-certificates \ 5 - libgtk-3-0t64 libdrm2 libgbm1 libxcb-dri3-0 \ 6 - inkscape \ 7 6 && rm -rf /var/lib/apt/lists/* 8 7 9 8 ARG INKSTITCH_VERSION=3.2.2 ··· 13 12 -O /tmp/inkstitch.tar.xz && \ 14 13 mkdir -p /opt/inkstitch && \ 15 14 tar xJf /tmp/inkstitch.tar.xz -C /opt/inkstitch && \ 16 - rm /tmp/inkstitch.tar.xz 15 + rm /tmp/inkstitch.tar.xz && \ 16 + rm -rf /opt/inkstitch/inkstitch/fonts \ 17 + /opt/inkstitch/inkstitch/bin/locales \ 18 + /opt/inkstitch/inkstitch/inx \ 19 + /opt/inkstitch/inkstitch/palettes \ 20 + /opt/inkstitch/inkstitch/tiles \ 21 + /opt/inkstitch/inkstitch/addons \ 22 + /opt/inkstitch/inkstitch/symbols \ 23 + /opt/inkstitch/inkstitch/dbus 24 + 25 + # Stage 2: runtime 26 + FROM python:3.11-slim 27 + 28 + RUN apt-get update && apt-get install -y --no-install-recommends \ 29 + inkscape libwayland-cursor0 libdrm2 libgbm1 libxcb-dri3-0 \ 30 + && rm -rf /var/lib/apt/lists/* 31 + 32 + COPY --from=builder /opt/inkstitch /opt/inkstitch 17 33 ENV INKSTITCH_BIN=/opt/inkstitch/inkstitch/bin/inkstitch 18 34 19 35 WORKDIR /app 20 - COPY requirements.txt . 36 + COPY worker/requirements.txt . 21 37 RUN pip install --no-cache-dir -r requirements.txt 22 38 23 - COPY pipeline/ pipeline/ 24 - COPY worker.py . 39 + COPY worker/pipeline/ pipeline/ 40 + COPY worker/worker.py . 25 41 26 - ENV RESULTS_DIR=/data/results 27 - CMD ["rq", "worker", "patches", "--url", "redis://redis:6379"] 42 + CMD ["python", "worker.py"]
+6 -12
docker-compose.yml
··· 5 5 - "6379:6379" 6 6 7 7 api: 8 - build: . 8 + build: 9 + context: . 10 + dockerfile: Dockerfile.web 9 11 ports: 10 12 - "8000:8000" 11 13 environment: 12 14 - REDIS_URL=redis://redis:6379 13 - - RESULTS_DIR=/data/results 14 - volumes: 15 - - results:/data/results 16 15 depends_on: 17 16 - redis 18 17 19 18 worker: 20 - build: . 21 - command: rq worker patches --url redis://redis:6379 19 + build: 20 + context: . 21 + dockerfile: Dockerfile.worker 22 22 environment: 23 23 - REDIS_URL=redis://redis:6379 24 - - RESULTS_DIR=/data/results 25 - volumes: 26 - - results:/data/results 27 24 depends_on: 28 25 - redis 29 - 30 - volumes: 31 - results:
-88
server.py
··· 1 - """ 2 - pic-to-patch API server. 3 - 4 - POST /patch — upload image, returns {job_id} 5 - GET /jobs/{job_id} — returns {status, result_url?} 6 - GET /jobs/{job_id}/result — returns the PNG directly 7 - 8 - Jobs are queued via Redis (rq). Results are stored in Redis as blobs. 9 - """ 10 - 11 - import os 12 - import uuid 13 - 14 - from fastapi import FastAPI, UploadFile, HTTPException 15 - from fastapi.responses import Response 16 - import redis 17 - import rq 18 - 19 - REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379") 20 - RESULT_TTL = int(os.environ.get("RESULT_TTL", "3600")) 21 - 22 - app = FastAPI(title="pic-to-patch", version="0.2.0") 23 - redis_conn = redis.from_url(REDIS_URL) 24 - queue = rq.Queue("patches", connection=redis_conn) 25 - 26 - 27 - @app.post("/patch") 28 - async def create_patch( 29 - file: UploadFile, 30 - border_color: str = "#0a0a14", 31 - color_precision: int = 8, 32 - postprocess: bool = True, 33 - ): 34 - job_id = str(uuid.uuid4()) 35 - input_bytes = await file.read() 36 - ext = (file.filename or "input.png").rsplit(".", 1)[-1] or "png" 37 - 38 - redis_conn.setex(f"job:{job_id}:input", RESULT_TTL, input_bytes) 39 - redis_conn.setex(f"job:{job_id}:ext", RESULT_TTL, ext) 40 - redis_conn.setex(f"job:{job_id}:status", RESULT_TTL, "processing") 41 - 42 - queue.enqueue( 43 - "worker.run_pipeline", 44 - job_id=job_id, 45 - border_color=border_color, 46 - color_precision=color_precision, 47 - postprocess=postprocess, 48 - job_timeout=300, 49 - result_ttl=RESULT_TTL, 50 - ) 51 - 52 - return {"job_id": job_id} 53 - 54 - 55 - @app.get("/jobs/{job_id}") 56 - async def get_job(job_id: str): 57 - status = redis_conn.get(f"job:{job_id}:status") 58 - if status is None: 59 - raise HTTPException(404, "job not found") 60 - 61 - status = status.decode() 62 - resp = {"status": status} 63 - 64 - if status == "complete": 65 - resp["result_url"] = f"/jobs/{job_id}/result" 66 - elif status == "failed": 67 - error = redis_conn.get(f"job:{job_id}:error") 68 - resp["error"] = error.decode() if error else "unknown error" 69 - 70 - return resp 71 - 72 - 73 - @app.get("/jobs/{job_id}/result") 74 - async def get_result(job_id: str): 75 - data = redis_conn.get(f"job:{job_id}:result") 76 - if data is None: 77 - raise HTTPException(404, "result not ready") 78 - return Response(content=data, media_type="image/png", 79 - headers={"Content-Disposition": f'inline; filename="{job_id}.png"'}) 80 - 81 - 82 - @app.get("/health") 83 - async def health(): 84 - try: 85 - redis_conn.ping() 86 - except Exception: 87 - raise HTTPException(503, "redis unavailable") 88 - return {"status": "ok"}
+13
web/go.mod
··· 1 + module pic-to-patch-web 2 + 3 + go 1.24 4 + 5 + require ( 6 + github.com/google/uuid v1.6.0 7 + github.com/redis/go-redis/v9 v9.19.0 8 + ) 9 + 10 + require ( 11 + github.com/cespare/xxhash/v2 v2.3.0 // indirect 12 + go.uber.org/atomic v1.11.0 // indirect 13 + )
+24
web/go.sum
··· 1 + github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= 2 + github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= 3 + github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= 4 + github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= 5 + github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= 6 + github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= 7 + github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 8 + github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 9 + github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= 10 + github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= 11 + github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE= 12 + github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= 13 + github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 14 + github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 15 + github.com/redis/go-redis/v9 v9.19.0 h1:XPVaaPSnG6RhYf7p+rmSa9zZfeVAnWsH5h3lxthOm/k= 16 + github.com/redis/go-redis/v9 v9.19.0/go.mod h1:v/M13XI1PVCDcm01VtPFOADfZtHf8YW3baQf57KlIkA= 17 + github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q= 18 + github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= 19 + github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= 20 + github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= 21 + go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= 22 + go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= 23 + golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc= 24 + golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+184
web/main.go
··· 1 + package main 2 + 3 + import ( 4 + "context" 5 + "encoding/json" 6 + "fmt" 7 + "io" 8 + "log" 9 + "net/http" 10 + "os" 11 + "strconv" 12 + "strings" 13 + "time" 14 + 15 + "github.com/google/uuid" 16 + "github.com/redis/go-redis/v9" 17 + ) 18 + 19 + var ( 20 + rdb *redis.Client 21 + resultTTL time.Duration 22 + ) 23 + 24 + type QueuePayload struct { 25 + JobID string `json:"job_id"` 26 + BorderColor string `json:"border_color"` 27 + ColorPrecision int `json:"color_precision"` 28 + Postprocess bool `json:"postprocess"` 29 + } 30 + 31 + func main() { 32 + redisURL := env("REDIS_URL", "redis://localhost:6379") 33 + ttlSec, _ := strconv.Atoi(env("RESULT_TTL", "3600")) 34 + resultTTL = time.Duration(ttlSec) * time.Second 35 + 36 + opt, err := redis.ParseURL(redisURL) 37 + if err != nil { 38 + log.Fatalf("bad REDIS_URL: %v", err) 39 + } 40 + rdb = redis.NewClient(opt) 41 + 42 + mux := http.NewServeMux() 43 + mux.HandleFunc("POST /patch", handleCreatePatch) 44 + mux.HandleFunc("GET /jobs/{job_id}", handleGetJob) 45 + mux.HandleFunc("GET /jobs/{job_id}/result", handleGetResult) 46 + mux.HandleFunc("GET /health", handleHealth) 47 + 48 + log.Println("listening on :8000") 49 + log.Fatal(http.ListenAndServe(":8000", mux)) 50 + } 51 + 52 + func handleCreatePatch(w http.ResponseWriter, r *http.Request) { 53 + file, header, err := r.FormFile("file") 54 + if err != nil { 55 + http.Error(w, "missing file", http.StatusBadRequest) 56 + return 57 + } 58 + defer file.Close() 59 + 60 + inputBytes, err := io.ReadAll(file) 61 + if err != nil { 62 + http.Error(w, "failed to read file", http.StatusInternalServerError) 63 + return 64 + } 65 + 66 + filename := "input.png" 67 + if header.Filename != "" { 68 + filename = header.Filename 69 + } 70 + ext := "png" 71 + if i := strings.LastIndex(filename, "."); i >= 0 && i+1 < len(filename) { 72 + ext = filename[i+1:] 73 + } 74 + 75 + borderColor := r.FormValue("border_color") 76 + if borderColor == "" { 77 + borderColor = "#0a0a14" 78 + } 79 + 80 + colorPrecision := 8 81 + if v := r.FormValue("color_precision"); v != "" { 82 + colorPrecision, _ = strconv.Atoi(v) 83 + } 84 + 85 + postprocess := true 86 + if v := r.FormValue("postprocess"); v != "" { 87 + postprocess = v != "false" && v != "0" && v != "no" 88 + } 89 + 90 + jobID := uuid.New().String() 91 + ctx := context.Background() 92 + 93 + pipe := rdb.Pipeline() 94 + pipe.SetEx(ctx, fmt.Sprintf("job:%s:input", jobID), inputBytes, resultTTL) 95 + pipe.SetEx(ctx, fmt.Sprintf("job:%s:ext", jobID), ext, resultTTL) 96 + pipe.SetEx(ctx, fmt.Sprintf("job:%s:status", jobID), "processing", resultTTL) 97 + if _, err := pipe.Exec(ctx); err != nil { 98 + http.Error(w, "redis error", http.StatusInternalServerError) 99 + return 100 + } 101 + 102 + payload := QueuePayload{ 103 + JobID: jobID, 104 + BorderColor: borderColor, 105 + ColorPrecision: colorPrecision, 106 + Postprocess: postprocess, 107 + } 108 + payloadJSON, _ := json.Marshal(payload) 109 + 110 + if err := rdb.LPush(ctx, "patches", payloadJSON).Err(); err != nil { 111 + http.Error(w, "queue error", http.StatusInternalServerError) 112 + return 113 + } 114 + 115 + writeJSON(w, http.StatusOK, map[string]string{"job_id": jobID}) 116 + } 117 + 118 + func handleGetJob(w http.ResponseWriter, r *http.Request) { 119 + jobID := r.PathValue("job_id") 120 + ctx := context.Background() 121 + 122 + status, err := rdb.Get(ctx, fmt.Sprintf("job:%s:status", jobID)).Result() 123 + if err == redis.Nil { 124 + http.Error(w, `{"detail":"job not found"}`, http.StatusNotFound) 125 + return 126 + } else if err != nil { 127 + http.Error(w, "redis error", http.StatusInternalServerError) 128 + return 129 + } 130 + 131 + resp := map[string]interface{}{"status": status} 132 + 133 + if status == "complete" { 134 + resp["result_url"] = fmt.Sprintf("/jobs/%s/result", jobID) 135 + } else if status == "failed" { 136 + errMsg, _ := rdb.Get(ctx, fmt.Sprintf("job:%s:error", jobID)).Result() 137 + if errMsg == "" { 138 + errMsg = "unknown error" 139 + } 140 + resp["error"] = errMsg 141 + } 142 + 143 + writeJSON(w, http.StatusOK, resp) 144 + } 145 + 146 + func handleGetResult(w http.ResponseWriter, r *http.Request) { 147 + jobID := r.PathValue("job_id") 148 + ctx := context.Background() 149 + 150 + data, err := rdb.Get(ctx, fmt.Sprintf("job:%s:result", jobID)).Bytes() 151 + if err == redis.Nil { 152 + http.Error(w, `{"detail":"result not ready"}`, http.StatusNotFound) 153 + return 154 + } else if err != nil { 155 + http.Error(w, "redis error", http.StatusInternalServerError) 156 + return 157 + } 158 + 159 + w.Header().Set("Content-Type", "image/png") 160 + w.Header().Set("Content-Disposition", fmt.Sprintf(`inline; filename="%s.png"`, jobID)) 161 + w.Write(data) 162 + } 163 + 164 + func handleHealth(w http.ResponseWriter, r *http.Request) { 165 + ctx := context.Background() 166 + if err := rdb.Ping(ctx).Err(); err != nil { 167 + http.Error(w, `{"detail":"redis unavailable"}`, http.StatusServiceUnavailable) 168 + return 169 + } 170 + writeJSON(w, http.StatusOK, map[string]string{"status": "ok"}) 171 + } 172 + 173 + func writeJSON(w http.ResponseWriter, status int, v interface{}) { 174 + w.Header().Set("Content-Type", "application/json") 175 + w.WriteHeader(status) 176 + json.NewEncoder(w).Encode(v) 177 + } 178 + 179 + func env(key, fallback string) string { 180 + if v := os.Getenv(key); v != "" { 181 + return v 182 + } 183 + return fallback 184 + }
-53
worker.py
··· 1 - """ 2 - pic-to-patch worker. Picks jobs from Redis queue and runs the pipeline. 3 - 4 - Run with: rq worker patches --url redis://localhost:6379 5 - """ 6 - 7 - import os 8 - import tempfile 9 - from pathlib import Path 10 - 11 - import redis 12 - 13 - from pipeline.convert import convert, convert_svg 14 - 15 - REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379") 16 - RESULT_TTL = int(os.environ.get("RESULT_TTL", "3600")) 17 - redis_conn = redis.from_url(REDIS_URL) 18 - 19 - 20 - def run_pipeline(job_id, border_color="#0a0a14", 21 - color_precision=8, postprocess=True): 22 - try: 23 - input_bytes = redis_conn.get(f"job:{job_id}:input") 24 - ext = (redis_conn.get(f"job:{job_id}:ext") or b"png").decode() 25 - if not input_bytes: 26 - raise RuntimeError("input not found in redis") 27 - 28 - is_svg = ext.lower() == "svg" 29 - 30 - with tempfile.TemporaryDirectory(prefix="p2p_") as tmpdir: 31 - tmpdir = Path(tmpdir) 32 - input_path = tmpdir / f"input.{ext}" 33 - input_path.write_bytes(input_bytes) 34 - output_path = tmpdir / "patch.png" 35 - 36 - if is_svg: 37 - convert_svg(str(input_path), str(output_path), 38 - border_color=border_color, postprocess=postprocess) 39 - else: 40 - convert(str(input_path), str(output_path), 41 - border_color=border_color, color_precision=color_precision, 42 - postprocess=postprocess) 43 - 44 - result_bytes = output_path.read_bytes() 45 - 46 - redis_conn.setex(f"job:{job_id}:result", RESULT_TTL, result_bytes) 47 - redis_conn.setex(f"job:{job_id}:status", RESULT_TTL, "complete") 48 - redis_conn.delete(f"job:{job_id}:input") 49 - 50 - except Exception as e: 51 - redis_conn.setex(f"job:{job_id}:status", RESULT_TTL, "failed") 52 - redis_conn.setex(f"job:{job_id}:error", RESULT_TTL, str(e)) 53 - raise
worker/pipeline/__init__.py

This is a binary file and will not be displayed.

+116
worker/pipeline/convert.py
··· 1 + """ 2 + Core pipeline: image → embroidered patch PNG. 3 + 4 + This is the throwaway python version. The interface is: 5 + convert(input_path, output_path, border_color="#0a0a14", color_precision=8) 6 + 7 + That's it. Everything else is internal. 8 + """ 9 + 10 + import os 11 + import subprocess 12 + import tempfile 13 + from pathlib import Path 14 + 15 + from PIL import Image 16 + import vtracer 17 + 18 + from .svg2patch import add_inkstitch_params, render_inkstitch 19 + 20 + MAX_IMAGE_DIM = int(os.environ.get("MAX_IMAGE_DIM", "500")) 21 + 22 + 23 + def convert(input_path, output_path, border_color="#0a0a14", color_precision=8, postprocess=True): 24 + input_path = Path(input_path) 25 + output_path = Path(output_path) 26 + 27 + with tempfile.TemporaryDirectory(prefix="p2p_") as tmpdir: 28 + tmpdir = Path(tmpdir) 29 + 30 + resized = tmpdir / "resized.png" 31 + _resize(input_path, resized) 32 + 33 + vtracer_svg = tmpdir / "vectorized.svg" 34 + _vectorize(resized, vtracer_svg, color_precision) 35 + 36 + embroidery_svg = tmpdir / "embroidery.svg" 37 + add_inkstitch_params(vtracer_svg, embroidery_svg, border_color=border_color) 38 + 39 + stitch_png = tmpdir / "stitch.png" 40 + if not render_inkstitch(embroidery_svg, stitch_png): 41 + raise RuntimeError("inkstitch render failed") 42 + 43 + if postprocess: 44 + _postprocess(stitch_png, output_path) 45 + else: 46 + import shutil 47 + shutil.copy2(stitch_png, output_path) 48 + 49 + return output_path 50 + 51 + 52 + def convert_svg(input_svg, output_path, border_color="#0a0a14", postprocess=True): 53 + input_svg = Path(input_svg) 54 + output_path = Path(output_path) 55 + 56 + with tempfile.TemporaryDirectory(prefix="p2p_") as tmpdir: 57 + tmpdir = Path(tmpdir) 58 + 59 + embroidery_svg = tmpdir / "embroidery.svg" 60 + add_inkstitch_params(input_svg, embroidery_svg, border_color=border_color) 61 + 62 + stitch_png = tmpdir / "stitch.png" 63 + if not render_inkstitch(embroidery_svg, stitch_png): 64 + raise RuntimeError("inkstitch render failed") 65 + 66 + if postprocess: 67 + _postprocess(stitch_png, output_path) 68 + else: 69 + import shutil 70 + shutil.copy2(stitch_png, output_path) 71 + 72 + return output_path 73 + 74 + 75 + def _resize(input_path, output_path): 76 + img = Image.open(input_path) 77 + if img.mode == "P": 78 + img = img.convert("RGBA") 79 + w, h = img.size 80 + if max(w, h) > MAX_IMAGE_DIM: 81 + ratio = MAX_IMAGE_DIM / max(w, h) 82 + img = img.resize((int(w * ratio), int(h * ratio)), Image.Resampling.LANCZOS) 83 + img.save(str(output_path)) 84 + 85 + 86 + def _vectorize(input_image, output_svg, color_precision=8): 87 + vtracer.convert_image_to_svg_py( 88 + image_path=str(input_image), 89 + out_path=str(output_svg), 90 + colormode="color", 91 + hierarchical="stacked", 92 + filter_speckle=4, 93 + color_precision=color_precision, 94 + corner_threshold=60, 95 + length_threshold=4, 96 + splice_threshold=45, 97 + ) 98 + if not output_svg.exists(): 99 + raise RuntimeError("vtracer produced no output") 100 + 101 + 102 + def _postprocess(stitch_png, output_path): 103 + script = Path(__file__).parent / "photorealistic.py" 104 + if not script.exists(): 105 + import shutil 106 + shutil.copy2(stitch_png, output_path) 107 + return 108 + 109 + import sys 110 + result = subprocess.run( 111 + [sys.executable, str(script), str(stitch_png), str(output_path)], 112 + capture_output=True, text=True, timeout=120, 113 + ) 114 + if result.returncode != 0 or not output_path.exists(): 115 + import shutil 116 + shutil.copy2(stitch_png, output_path)
+764
worker/pipeline/photorealistic.py
··· 1 + #!/usr/bin/env python3 2 + """ 3 + Advanced post-processing pipeline that makes inkstitch renders look like 4 + photographs of real embroidered patches. 5 + 6 + Pipeline: 7 + 1. Extract patch mask from white background 8 + 2. Compute per-pixel normal map from stitch texture gradients 9 + 3. Estimate local thread direction via structure tensor 10 + 4. Apply Blinn-Phong shading with Kajiya-Kay anisotropic specular 11 + 5. Add per-thread micro-highlight variation 12 + 6. Generate felt/twill backing fabric texture 13 + 7. Create realistic drop shadow (contact + cast) 14 + 8. Simulate merrow/overlock edge border 15 + 9. Add patch thickness bevel (embossed raised edge) 16 + 10. Composite everything 17 + 11. Photographic finishing: DOF, color grade, vignette, grain 18 + """ 19 + 20 + import sys 21 + import numpy as np 22 + from pathlib import Path 23 + from PIL import Image, ImageFilter, ImageEnhance, ImageDraw 24 + from scipy import ndimage 25 + from scipy.ndimage import gaussian_filter, uniform_filter 26 + 27 + 28 + # ───────────────────────────────────────────── 29 + # 1. Mask extraction 30 + # ───────────────────────────────────────────── 31 + 32 + def extract_patch_mask(img_arr, threshold=235): 33 + """ 34 + Extract a mask of the embroidered region vs white background. 35 + Returns float mask [0, 1] with anti-aliased soft edges. 36 + """ 37 + is_bg = np.all(img_arr[:, :, :3] > threshold, axis=2) 38 + mask = (~is_bg).astype(np.float32) 39 + 40 + # Close small holes inside the patch 41 + mask = ndimage.binary_closing(mask > 0.5, iterations=3).astype(np.float32) 42 + # Dilate slightly to catch anti-aliased edge pixels 43 + mask = ndimage.binary_dilation(mask > 0.5, iterations=1).astype(np.float32) 44 + 45 + # Soft edge via gaussian blur 46 + mask_soft = gaussian_filter(mask, sigma=1.2) 47 + return np.clip(mask_soft, 0, 1) 48 + 49 + 50 + # ───────────────────────────────────────────── 51 + # 2. Normal map computation 52 + # ───────────────────────────────────────────── 53 + 54 + def compute_normal_map(img_arr, mask, strength=2.0): 55 + """ 56 + Derive per-pixel surface normals from stitch texture luminance gradients. 57 + Multi-scale: fine captures thread ridges, medium captures stitch rows. 58 + """ 59 + lum = (0.299 * img_arr[:, :, 0] + 60 + 0.587 * img_arr[:, :, 1] + 61 + 0.114 * img_arr[:, :, 2]).astype(np.float64) 62 + 63 + # Fine scale - individual thread ridges 64 + gx_fine = ndimage.sobel(lum, axis=1) 65 + gy_fine = ndimage.sobel(lum, axis=0) 66 + 67 + # Medium scale - stitch row structure 68 + lum_med = gaussian_filter(lum, sigma=1.2) 69 + gx_med = ndimage.sobel(lum_med, axis=1) 70 + gy_med = ndimage.sobel(lum_med, axis=0) 71 + 72 + # Coarse scale - broad curvature 73 + lum_coarse = gaussian_filter(lum, sigma=3.0) 74 + gx_coarse = ndimage.sobel(lum_coarse, axis=1) 75 + gy_coarse = ndimage.sobel(lum_coarse, axis=0) 76 + 77 + gx = 0.5 * gx_fine + 0.35 * gx_med + 0.15 * gx_coarse 78 + gy = 0.5 * gy_fine + 0.35 * gy_med + 0.15 * gy_coarse 79 + 80 + gx *= strength 81 + gy *= strength 82 + 83 + h, w = lum.shape 84 + normals = np.zeros((h, w, 3), dtype=np.float64) 85 + normals[:, :, 0] = -gx 86 + normals[:, :, 1] = -gy 87 + normals[:, :, 2] = 255.0 88 + 89 + length = np.sqrt(np.sum(normals ** 2, axis=2, keepdims=True)) 90 + normals /= np.maximum(length, 1e-8) 91 + normals *= mask[:, :, np.newaxis] 92 + 93 + return normals 94 + 95 + 96 + # ───────────────────────────────────────────── 97 + # 3. Thread direction estimation 98 + # ───────────────────────────────────────────── 99 + 100 + def estimate_thread_direction(img_arr, mask, block_size=12): 101 + """ 102 + Estimate local thread direction via structure tensor analysis. 103 + Returns angle map in radians for anisotropic specular. 104 + """ 105 + lum = (0.299 * img_arr[:, :, 0] + 106 + 0.587 * img_arr[:, :, 1] + 107 + 0.114 * img_arr[:, :, 2]).astype(np.float64) 108 + 109 + gx = ndimage.sobel(lum, axis=1) 110 + gy = ndimage.sobel(lum, axis=0) 111 + 112 + sigma = block_size / 2.0 113 + Jxx = gaussian_filter(gx * gx, sigma=sigma) 114 + Jxy = gaussian_filter(gx * gy, sigma=sigma) 115 + Jyy = gaussian_filter(gy * gy, sigma=sigma) 116 + 117 + # Dominant gradient direction 118 + angle = 0.5 * np.arctan2(2.0 * Jxy, Jxx - Jyy + 1e-10) 119 + # Thread runs perpendicular to gradient 120 + thread_angle = angle + np.pi / 2.0 121 + 122 + # Also compute anisotropy strength (how directional the texture is) 123 + trace = Jxx + Jyy + 1e-10 124 + diff = np.sqrt((Jxx - Jyy) ** 2 + 4 * Jxy ** 2) 125 + anisotropy = diff / trace # 0 = isotropic, 1 = perfectly directional 126 + 127 + return thread_angle, anisotropy 128 + 129 + 130 + # ───────────────────────────────────────────── 131 + # 4. Blinn-Phong + anisotropic shading 132 + # ───────────────────────────────────────────── 133 + 134 + def apply_lighting(img_arr, normals, mask, thread_angle, anisotropy, 135 + light_dir=(0.25, -0.35, 0.90), 136 + ambient=0.55, diffuse_strength=0.35, 137 + specular_strength=0.12, shininess=18.0, 138 + aniso_strength=0.08, aniso_shininess=30.0): 139 + """ 140 + Apply Blinn-Phong shading with Kajiya-Kay anisotropic specular. 141 + 142 + The key insight: we DON'T want to darken the patch too much. Real embroidery 143 + is bright and saturated under normal lighting. The shading should add 144 + subtle variation and highlights, not dramatically change overall brightness. 145 + 146 + We use a "detail lighting" approach: 147 + final = original * (ambient + diffuse) + specular_white 148 + where ambient+diffuse averages close to 1.0 across the image. 149 + """ 150 + result = img_arr.astype(np.float64).copy() 151 + 152 + light = np.array(light_dir, dtype=np.float64) 153 + light /= np.linalg.norm(light) 154 + 155 + view = np.array([0.0, 0.0, 1.0]) 156 + half_vec = light + view 157 + half_vec /= np.linalg.norm(half_vec) 158 + 159 + # --- Diffuse --- 160 + n_dot_l = (normals[:, :, 0] * light[0] + 161 + normals[:, :, 1] * light[1] + 162 + normals[:, :, 2] * light[2]) 163 + n_dot_l = np.clip(n_dot_l, 0, 1) 164 + 165 + # Wrap diffuse to soften shadows (half-lambert) 166 + diffuse = diffuse_strength * (n_dot_l * 0.5 + 0.5) 167 + 168 + # --- Isotropic specular --- 169 + n_dot_h = (normals[:, :, 0] * half_vec[0] + 170 + normals[:, :, 1] * half_vec[1] + 171 + normals[:, :, 2] * half_vec[2]) 172 + n_dot_h = np.clip(n_dot_h, 0, 1) 173 + specular_iso = specular_strength * np.power(n_dot_h, shininess) 174 + 175 + # --- Anisotropic specular (Kajiya-Kay) --- 176 + tangent_x = np.cos(thread_angle) 177 + tangent_y = np.sin(thread_angle) 178 + 179 + t_dot_h = tangent_x * half_vec[0] + tangent_y * half_vec[1] 180 + sin_th = np.sqrt(np.clip(1.0 - t_dot_h ** 2, 0, 1)) 181 + # Modulate by local anisotropy - only apply aniso spec where texture is directional 182 + specular_aniso = aniso_strength * np.power(sin_th, aniso_shininess) * anisotropy 183 + 184 + # --- Secondary broad specular (satin sheen on lighter regions) --- 185 + # Lighter areas (white text, bright embroidery) have more specular 186 + lum = np.mean(result[:, :, :3], axis=2) / 255.0 187 + brightness_boost = np.clip(lum - 0.4, 0, 0.6) / 0.6 # ramp from 0.4 to 1.0 188 + spec_broad = 0.06 * np.power(n_dot_h, 8.0) * brightness_boost 189 + 190 + # --- Combine --- 191 + # Multiplicative lighting (affects color) 192 + color_light = ambient + diffuse 193 + # Ensure average is near 1.0 to preserve original brightness 194 + color_light = np.clip(color_light, 0.35, 1.5) 195 + 196 + # Additive specular (white highlight) 197 + spec_total = (specular_iso + specular_aniso + spec_broad) * mask 198 + 199 + # Apply 200 + for c in range(3): 201 + channel = result[:, :, c] 202 + lit = channel * color_light + spec_total * 220.0 203 + result[:, :, c] = channel * (1.0 - mask) + lit * mask 204 + 205 + return np.clip(result, 0, 255) 206 + 207 + 208 + # ───────────────────────────────────────────── 209 + # 5. Per-thread micro-highlights 210 + # ───────────────────────────────────────────── 211 + 212 + def compute_ambient_occlusion(img_arr, mask, radius=2.0, strength=0.15): 213 + """ 214 + Approximate screen-space ambient occlusion from stitch texture. 215 + Stitch valleys (between thread rows) are slightly darker. 216 + Computed by comparing local brightness to neighborhood average. 217 + """ 218 + lum = np.mean(img_arr[:, :, :3], axis=2) 219 + local_avg = gaussian_filter(lum, sigma=radius) 220 + ao = np.clip((local_avg - lum) / (local_avg + 1e-8), 0, 1) * strength 221 + ao *= mask 222 + return ao 223 + 224 + 225 + def add_thread_microhighlights(img_arr, normals, mask, thread_angle, intensity=0.05): 226 + """ 227 + Individual threads catch light slightly differently. 228 + Creates fine-grained brightness variation aligned with thread direction. 229 + Also adds subtle shimmer for thread luster. 230 + """ 231 + h, w = img_arr.shape[:2] 232 + noise = np.random.normal(0, 1, (h, w)) 233 + 234 + # Directional blur kernels 235 + noise_along = gaussian_filter(noise, sigma=[0.3, 3.0]) 236 + noise_across = gaussian_filter(noise, sigma=[3.0, 0.3]) 237 + 238 + # Blend using thread angle 239 + cos_a = np.abs(np.cos(thread_angle)) 240 + sin_a = np.abs(np.sin(thread_angle)) 241 + total = cos_a + sin_a + 1e-8 242 + directional = (noise_along * cos_a + noise_across * sin_a) / total 243 + 244 + highlight = directional * intensity * mask 245 + result = img_arr.copy() 246 + for c in range(3): 247 + result[:, :, c] *= (1.0 + highlight) 248 + 249 + # Fine per-pixel shimmer for thread luster 250 + shimmer = np.random.normal(0, 0.015, (h, w)) 251 + shimmer = gaussian_filter(shimmer, sigma=0.3) 252 + for c in range(3): 253 + result[:, :, c] *= (1.0 + shimmer * mask) 254 + 255 + return np.clip(result, 0, 255) 256 + 257 + 258 + # ───────────────────────────────────────────── 259 + # 6. Fabric backing texture 260 + # ───────────────────────────────────────────── 261 + 262 + def generate_fabric_texture(width, height, color=(35, 35, 40), weave_scale=3): 263 + """ 264 + Generate realistic felt/twill fabric backing. 265 + Multi-octave noise with diagonal weave pattern. 266 + """ 267 + arr = np.full((height, width, 3), color, dtype=np.float64) 268 + 269 + y_idx = np.arange(height)[:, None] 270 + x_idx = np.arange(width)[None, :] 271 + 272 + # Twill diagonal weave 273 + twill1 = np.sin(2 * np.pi * (x_idx + y_idx) / weave_scale) * 0.035 274 + twill2 = np.sin(2 * np.pi * (x_idx - y_idx) / (weave_scale * 1.3)) * 0.02 275 + twill3 = np.sin(2 * np.pi * (x_idx * 0.7 + y_idx * 1.3) / (weave_scale * 2)) * 0.015 276 + arr *= (1.0 + twill1 + twill2 + twill3)[:, :, np.newaxis] 277 + 278 + # Fine fiber noise 279 + noise_fine = np.random.normal(0, 2.0, (height, width, 3)) 280 + # Medium texture clumps 281 + noise_med_small = np.random.normal(0, 1.2, (height // 2 + 1, width // 2 + 1, 3)) 282 + noise_med = np.repeat(np.repeat(noise_med_small, 2, axis=0), 2, axis=1)[:height, :width, :] 283 + noise_med = gaussian_filter(noise_med, sigma=1.0) 284 + # Coarse color drift 285 + noise_coarse_small = np.random.normal(0, 0.8, (height // 6 + 1, width // 6 + 1, 3)) 286 + noise_coarse = np.repeat(np.repeat(noise_coarse_small, 6, axis=0), 6, axis=1)[:height, :width, :] 287 + noise_coarse = gaussian_filter(noise_coarse, sigma=3.0) 288 + 289 + arr += noise_fine + noise_med + noise_coarse 290 + 291 + # Subtle large-scale brightness variation (fabric isn't perfectly uniform) 292 + var_small = np.random.normal(0, 0.015, (height // 12 + 1, width // 12 + 1)) 293 + variation = np.repeat(np.repeat(var_small, 12, axis=0), 12, axis=1)[:height, :width] 294 + variation = gaussian_filter(variation, sigma=6.0) 295 + arr *= (1.0 + variation[:, :, np.newaxis]) 296 + 297 + return np.clip(arr, 0, 255).astype(np.uint8) 298 + 299 + 300 + # ───────────────────────────────────────────── 301 + # 7. Shadow 302 + # ───────────────────────────────────────────── 303 + 304 + def create_patch_shadow(mask, offset=(5, 6), blur_radius=10, opacity=0.55): 305 + """ 306 + Realistic two-layer shadow: 307 + - Contact shadow: tight, dark, right at the edge 308 + - Cast shadow: offset, soft, diffuse 309 + """ 310 + h, w = mask.shape 311 + 312 + # Cast shadow: shift the mask by offset and blur 313 + shifted = ndimage.shift(mask, (offset[1], offset[0]), order=1, mode='constant', cval=0) 314 + cast = gaussian_filter(shifted, sigma=blur_radius) * opacity 315 + 316 + # Contact shadow: unshifted tight edge glow 317 + # Expand mask slightly, subtract original, blur tightly 318 + dilated = gaussian_filter(mask, sigma=2.0) 319 + contact_ring = np.clip(dilated - mask * 0.9, 0, 1) 320 + contact = gaussian_filter(contact_ring, sigma=2.5) * 0.4 321 + 322 + # Combine 323 + shadow = np.maximum(cast, contact) 324 + # Don't darken the patch interior 325 + shadow *= np.clip(1.0 - mask, 0, 1) 326 + 327 + return np.clip(shadow, 0, 1) 328 + 329 + 330 + # ───────────────────────────────────────────── 331 + # 8. Merrow edge 332 + # ───────────────────────────────────────────── 333 + 334 + def create_merrow_edge(mask, thickness=3): 335 + """ 336 + Simulate overlock/merrow stitch border around patch edge. 337 + Returns (edge_mask, edge_normals) for 3D stitched border look. 338 + """ 339 + h, w = mask.shape 340 + 341 + # Get edge band via morphological gradient 342 + hard_mask = (mask > 0.5).astype(np.float32) 343 + dilated = ndimage.binary_dilation(hard_mask > 0.5, iterations=thickness).astype(np.float32) 344 + eroded = ndimage.binary_erosion(hard_mask > 0.5, iterations=max(1, thickness // 2)).astype(np.float32) 345 + edge_band = np.clip(dilated - eroded, 0, 1) 346 + 347 + # Anti-alias the edge band 348 + edge_band = gaussian_filter(edge_band, sigma=0.6) 349 + 350 + # Add stitch texture pattern along the edge 351 + y_idx = np.arange(h)[:, None].astype(np.float64) 352 + x_idx = np.arange(w)[None, :].astype(np.float64) 353 + 354 + # Distance from center for radial stitch direction 355 + cy, cx = h / 2.0, w / 2.0 356 + dy = y_idx - cy 357 + dx = x_idx - cx 358 + angle = np.arctan2(dy, dx) 359 + 360 + # Stitch pattern follows the edge circumferentially 361 + # Use angle to create perpendicular stitches 362 + dist = np.sqrt(dx ** 2 + dy ** 2) 363 + stitch_freq = dist * 0.15 # scale frequency by radius 364 + stitch_pattern = np.sin(stitch_freq + angle * 25) * 0.12 + 0.88 365 + stitch_pattern2 = np.cos(stitch_freq * 1.7 + angle * 18) * 0.06 + 0.94 366 + 367 + edge_textured = edge_band * stitch_pattern * stitch_pattern2 368 + 369 + return edge_textured, edge_band 370 + 371 + 372 + # ───────────────────────────────────────────── 373 + # 9. Patch thickness bevel 374 + # ───────────────────────────────────────────── 375 + 376 + def create_edge_bevel(mask, bevel_width=6, light_dir=(0.3, -0.4)): 377 + """ 378 + Create a bevel/emboss effect at the patch edge to simulate thickness. 379 + The patch is raised ~1-2mm above the fabric, creating a lit top edge 380 + and shadowed bottom edge. 381 + """ 382 + # Compute distance from edge (inward) 383 + hard_mask = (mask > 0.5).astype(np.float32) 384 + dist = ndimage.distance_transform_edt(hard_mask) 385 + dist_outside = ndimage.distance_transform_edt(1 - hard_mask) 386 + 387 + # Bevel height profile: ramps up at edge, flat in interior 388 + bevel_height = np.clip(dist / bevel_width, 0, 1) 389 + # Also slight ramp outside for the outer edge 390 + bevel_height_out = np.clip(1.0 - dist_outside / (bevel_width * 0.5), 0, 1) * (1 - hard_mask) 391 + 392 + height = bevel_height + bevel_height_out 393 + 394 + # Compute lighting from height map 395 + gx = ndimage.sobel(height, axis=1) 396 + gy = ndimage.sobel(height, axis=0) 397 + 398 + # Directional lighting 399 + lx, ly = light_dir 400 + bevel_light = -(gx * lx + gy * ly) 401 + 402 + # Normalize to [-1, 1] range 403 + max_val = max(np.abs(bevel_light).max(), 1e-8) 404 + bevel_light = bevel_light / max_val 405 + 406 + # Only apply near edges 407 + edge_proximity = np.clip(1.0 - dist / (bevel_width * 1.5), 0, 1) * hard_mask 408 + edge_proximity += np.clip(1.0 - dist_outside / (bevel_width * 0.8), 0, 1) * (1 - hard_mask) 409 + 410 + bevel_light *= edge_proximity 411 + 412 + return bevel_light 413 + 414 + 415 + def create_inner_relief(img_arr, mask, light_dir=(0.3, -0.4), strength=0.08): 416 + """ 417 + Detect color boundaries within the patch (where different stitch sections 418 + meet) and add subtle height relief at those boundaries. In real embroidery, 419 + the NASA text sits slightly above the blue fill, the chevron overlaps, etc. 420 + """ 421 + h, w = img_arr.shape[:2] 422 + 423 + # Detect edges within the patch using color gradient magnitude 424 + # Use all 3 channels for better boundary detection 425 + edges = np.zeros((h, w), dtype=np.float64) 426 + for c in range(3): 427 + gx = ndimage.sobel(img_arr[:, :, c], axis=1) 428 + gy = ndimage.sobel(img_arr[:, :, c], axis=0) 429 + edges += np.sqrt(gx ** 2 + gy ** 2) 430 + edges /= 3.0 431 + 432 + # Threshold to find significant color boundaries (not just stitch texture) 433 + # Smooth to get section-level boundaries, not individual thread edges 434 + edges_smooth = gaussian_filter(edges, sigma=1.5) 435 + 436 + # Normalize 437 + edge_max = np.percentile(edges_smooth[mask > 0.5], 95) if np.any(mask > 0.5) else 1.0 438 + edges_norm = np.clip(edges_smooth / (edge_max + 1e-8), 0, 1) 439 + 440 + # Create height map: sections have flat heights, boundaries have transitions 441 + # Use edge magnitude as a proxy for height discontinuity 442 + height = gaussian_filter(edges_norm, sigma=2.0) * mask 443 + 444 + # Compute directional lighting on this height map 445 + gx = ndimage.sobel(height, axis=1) 446 + gy = ndimage.sobel(height, axis=0) 447 + lx, ly = light_dir 448 + relief = -(gx * lx + gy * ly) * strength * mask 449 + 450 + return relief 451 + 452 + 453 + # ───────────────────────────────────────────── 454 + # 10-11. Photographic effects 455 + # ───────────────────────────────────────────── 456 + 457 + def add_vignette(img_arr, strength=0.22, radius=0.65): 458 + """Photographic vignette - darkens corners.""" 459 + h, w = img_arr.shape[:2] 460 + y = np.linspace(-1, 1, h)[:, None] 461 + x = np.linspace(-1, 1, w)[None, :] 462 + dist = np.sqrt(x * x + y * y) 463 + vignette = 1.0 - strength * np.clip((dist - radius) / (1.4 - radius), 0, 1) ** 1.5 464 + return img_arr * vignette[:, :, np.newaxis] 465 + 466 + 467 + def add_film_grain(img_arr, strength=3.0): 468 + """Photographic film grain with realistic grain size.""" 469 + h, w = img_arr.shape[:2] 470 + # Luminance-dependent grain (stronger in shadows) 471 + lum = np.mean(img_arr[:, :, :3], axis=2) 472 + grain_strength = strength * (1.0 + 0.3 * (1.0 - lum / 255.0)) 473 + 474 + grain = np.random.normal(0, 1, (h, w)) * grain_strength 475 + grain = gaussian_filter(grain, sigma=0.4) 476 + 477 + result = img_arr + grain[:, :, np.newaxis] 478 + return np.clip(result, 0, 255) 479 + 480 + 481 + def add_depth_of_field(img_arr, mask, max_blur=1.3): 482 + """Subtle DOF: patch center sharp, frame edges soft.""" 483 + h, w = img_arr.shape[:2] 484 + 485 + y_coords, x_coords = np.where(mask > 0.5) 486 + if len(y_coords) == 0: 487 + return img_arr 488 + cx, cy = int(np.mean(x_coords)), int(np.mean(y_coords)) 489 + 490 + y = np.arange(h)[:, None] 491 + x = np.arange(w)[None, :] 492 + dist = np.sqrt(((x - cx) / w * 2) ** 2 + ((y - cy) / h * 2) ** 2) 493 + blur_t = np.clip((dist - 0.35) / 0.65, 0, 1) ** 1.5 494 + 495 + blurred = np.stack([ 496 + gaussian_filter(img_arr[:, :, c], sigma=max_blur) 497 + for c in range(img_arr.shape[2]) 498 + ], axis=2) 499 + 500 + blend = blur_t[:, :, np.newaxis] 501 + return img_arr * (1 - blend) + blurred * blend 502 + 503 + 504 + def color_grade(img_arr, warmth=0.02, contrast=1.05, saturation=1.10): 505 + """Subtle photographic color grading.""" 506 + result = img_arr.copy() 507 + 508 + # Warmth 509 + result[:, :, 0] *= (1.0 + warmth) 510 + result[:, :, 1] *= (1.0 + warmth * 0.2) 511 + result[:, :, 2] *= (1.0 - warmth * 0.4) 512 + 513 + # S-curve contrast (gentle) 514 + mid = 128.0 515 + result = mid + (result - mid) * contrast 516 + 517 + # Saturation boost (embroidery is vivid) 518 + gray = np.mean(result[:, :, :3], axis=2, keepdims=True) 519 + result[:, :, :3] = gray + (result[:, :, :3] - gray) * saturation 520 + 521 + return np.clip(result, 0, 255) 522 + 523 + 524 + # ───────────────────────────────────────────── 525 + # Main pipeline 526 + # ───────────────────────────────────────────── 527 + 528 + def postprocess_photorealistic(input_path, output_path, 529 + fabric_color=(35, 35, 40), 530 + padding=55, 531 + light_dir=(0.25, -0.35, 0.90)): 532 + """ 533 + Full photorealistic post-processing pipeline. 534 + """ 535 + print(f"Loading {input_path}...") 536 + stitch_img = Image.open(input_path) 537 + stitch_arr = np.array(stitch_img)[:, :, :3].astype(np.float64) 538 + h, w = stitch_arr.shape[:2] 539 + 540 + print(" [1/13] Extracting patch mask...") 541 + mask = extract_patch_mask(stitch_arr) 542 + 543 + print(" [2/13] Computing normal map...") 544 + normals = compute_normal_map(stitch_arr, mask, strength=2.0) 545 + 546 + print(" [3/13] Estimating thread directions...") 547 + thread_angle, anisotropy = estimate_thread_direction(stitch_arr, mask, block_size=12) 548 + 549 + print(" [4/13] Applying Blinn-Phong + anisotropic shading...") 550 + lit = apply_lighting(stitch_arr, normals, mask, thread_angle, anisotropy, 551 + light_dir=light_dir, 552 + ambient=0.55, diffuse_strength=0.35, 553 + specular_strength=0.12, shininess=18.0, 554 + aniso_strength=0.08, aniso_shininess=30.0) 555 + 556 + print(" [5/13] Computing ambient occlusion...") 557 + ao = compute_ambient_occlusion(stitch_arr, mask, radius=2.0, strength=0.12) 558 + # Apply AO: darken stitch valleys 559 + for c in range(3): 560 + lit[:, :, c] *= (1.0 - ao) 561 + lit = np.clip(lit, 0, 255) 562 + 563 + print(" [6/13] Adding per-thread micro-highlights...") 564 + lit = add_thread_microhighlights(lit, normals, mask, thread_angle, intensity=0.05) 565 + 566 + # --- Canvas setup --- 567 + canvas_w = w + padding * 2 568 + canvas_h = h + padding * 2 569 + 570 + print(f" [7/13] Generating fabric texture ({canvas_w}x{canvas_h})...") 571 + fabric = generate_fabric_texture(canvas_w, canvas_h, color=fabric_color) 572 + canvas = fabric.astype(np.float64) 573 + 574 + # --- Pad mask to canvas size --- 575 + mask_padded = np.zeros((canvas_h, canvas_w), dtype=np.float64) 576 + mask_padded[padding:padding + h, padding:padding + w] = mask 577 + 578 + # --- Shadow --- 579 + print(" [8/13] Creating drop shadow...") 580 + shadow = create_patch_shadow(mask_padded, offset=(5, 6), blur_radius=10, opacity=0.50) 581 + for c in range(3): 582 + canvas[:, :, c] *= (1.0 - shadow * 0.8) 583 + 584 + # --- Composite lit patch onto canvas --- 585 + print(" [9/13] Compositing patch onto fabric...") 586 + for c in range(3): 587 + patch_channel = lit[:, :, c] 588 + canvas_region = canvas[padding:padding + h, padding:padding + w, c] 589 + canvas[padding:padding + h, padding:padding + w, c] = ( 590 + canvas_region * (1 - mask) + patch_channel * mask 591 + ) 592 + 593 + # --- Edge bevel (thickness illusion) --- 594 + print(" [10/13] Adding edge bevel for 3D thickness...") 595 + bevel = create_edge_bevel(mask_padded, bevel_width=5, 596 + light_dir=(light_dir[0], light_dir[1])) 597 + # Apply bevel as brightness modulation 598 + bevel_intensity = 45.0 # How strong the bevel highlight/shadow is 599 + for c in range(3): 600 + canvas[:, :, c] += bevel * bevel_intensity 601 + canvas = np.clip(canvas, 0, 255) 602 + 603 + # --- Inner relief --- 604 + print(" [11/13] Adding inner relief at section boundaries...") 605 + inner_relief = create_inner_relief(stitch_arr, mask, 606 + light_dir=(light_dir[0], light_dir[1]), 607 + strength=0.07) 608 + relief_padded = np.zeros((canvas_h, canvas_w), dtype=np.float64) 609 + relief_padded[padding:padding + h, padding:padding + w] = inner_relief 610 + for c in range(3): 611 + canvas[:, :, c] *= (1.0 + relief_padded) 612 + canvas = np.clip(canvas, 0, 255) 613 + 614 + print(" [12/13] Adding merrow edge border...") 615 + merrow, merrow_band = create_merrow_edge(mask_padded, thickness=3) 616 + 617 + # Detect dominant edge color from the patch border pixels 618 + # Sample colors from the edge region of the lit patch 619 + edge_sample_mask = (mask > 0.3) & (mask < 0.95) 620 + if np.any(edge_sample_mask): 621 + edge_colors = lit[edge_sample_mask] 622 + avg_edge = np.mean(edge_colors, axis=0) 623 + # Make merrow edge slightly lighter than border 624 + edge_color = np.clip(avg_edge * 1.3 + 30, 0, 255) 625 + else: 626 + edge_color = np.array([160, 160, 165], dtype=np.float64) 627 + 628 + # Apply merrow edge with bump lighting for 3D thread appearance 629 + # Create mini-normal from merrow pattern for lit edge 630 + merrow_gx = ndimage.sobel(merrow, axis=1) 631 + merrow_gy = ndimage.sobel(merrow, axis=0) 632 + merrow_light = -(merrow_gx * light_dir[0] + merrow_gy * light_dir[1]) 633 + merrow_light = merrow_light / (np.abs(merrow_light).max() + 1e-8) * 0.3 634 + 635 + for c in range(3): 636 + edge_val = edge_color[c] * (1.0 + merrow_light) * merrow 637 + canvas[:, :, c] = canvas[:, :, c] * (1 - merrow_band * 0.65) + edge_val * 0.65 638 + canvas = np.clip(canvas, 0, 255) 639 + 640 + # --- Photographic finishing --- 641 + print(" [13/13] Photographic finishing (DOF, color, vignette, grain)...") 642 + canvas = add_depth_of_field(canvas, mask_padded, max_blur=1.2) 643 + canvas = color_grade(canvas, warmth=0.02, contrast=1.06, saturation=1.12) 644 + canvas = add_vignette(canvas, strength=0.20, radius=0.60) 645 + canvas = add_film_grain(canvas, strength=3.0) 646 + 647 + # --- Save --- 648 + result = np.clip(canvas, 0, 255).astype(np.uint8) 649 + output_img = Image.fromarray(result, "RGB") 650 + output_img.save(str(output_path), quality=95) 651 + print(f" Saved to {output_path}") 652 + 653 + return output_path 654 + 655 + 656 + def save_debug_stages(input_path, output_dir): 657 + """Save intermediate stages for inspection.""" 658 + output_dir = Path(output_dir) 659 + output_dir.mkdir(parents=True, exist_ok=True) 660 + 661 + stitch_img = Image.open(input_path) 662 + stitch_arr = np.array(stitch_img)[:, :, :3].astype(np.float64) 663 + h, w = stitch_arr.shape[:2] 664 + 665 + mask = extract_patch_mask(stitch_arr) 666 + normals = compute_normal_map(stitch_arr, mask, strength=2.0) 667 + thread_angle, anisotropy = estimate_thread_direction(stitch_arr, mask, block_size=12) 668 + 669 + # 1. Normal map visualization (standard purple/green/blue encoding) 670 + normal_vis = ((normals + 1.0) * 0.5 * 255).astype(np.uint8) 671 + Image.fromarray(normal_vis, "RGB").save(str(output_dir / "01_normal_map.png")) 672 + 673 + # 2. Mask 674 + mask_vis = (mask * 255).astype(np.uint8) 675 + Image.fromarray(mask_vis, "L").save(str(output_dir / "02_mask.png")) 676 + 677 + # 3. Thread direction + anisotropy 678 + thread_vis = np.zeros((h, w, 3), dtype=np.uint8) 679 + thread_vis[:, :, 0] = ((np.cos(thread_angle) + 1) * 0.5 * 255 * mask).astype(np.uint8) 680 + thread_vis[:, :, 1] = ((np.sin(thread_angle) + 1) * 0.5 * 255 * mask).astype(np.uint8) 681 + thread_vis[:, :, 2] = (anisotropy * 255 * mask).astype(np.uint8) 682 + Image.fromarray(thread_vis, "RGB").save(str(output_dir / "03_thread_direction.png")) 683 + 684 + # 4. Lit patch (after shading, before compositing) 685 + lit = apply_lighting(stitch_arr, normals, mask, thread_angle, anisotropy) 686 + lit = add_thread_microhighlights(lit, normals, mask, thread_angle) 687 + Image.fromarray(np.clip(lit, 0, 255).astype(np.uint8), "RGB").save( 688 + str(output_dir / "04_lit_patch.png")) 689 + 690 + # 5. Edge bevel visualization 691 + mask_padded = np.zeros((h + 100, w + 100), dtype=np.float64) 692 + mask_padded[50:50 + h, 50:50 + w] = mask 693 + bevel = create_edge_bevel(mask_padded, bevel_width=5) 694 + bevel_vis = ((bevel + 1) * 0.5 * 255).astype(np.uint8) 695 + Image.fromarray(bevel_vis, "L").save(str(output_dir / "05_edge_bevel.png")) 696 + 697 + # 6. Shadow 698 + shadow = create_patch_shadow(mask_padded) 699 + shadow_vis = (shadow * 255).astype(np.uint8) 700 + Image.fromarray(shadow_vis, "L").save(str(output_dir / "06_shadow.png")) 701 + 702 + print(f"Debug stages saved to {output_dir}/") 703 + 704 + 705 + def create_comparison(input_path, output_path, comparison_path): 706 + """Create side-by-side comparison image.""" 707 + original = Image.open(input_path) 708 + result = Image.open(output_path) 709 + 710 + # Make them the same height 711 + orig_w, orig_h = original.size 712 + res_w, res_h = result.size 713 + 714 + # Scale original to match result height 715 + scale = res_h / orig_h 716 + orig_scaled = original.resize((int(orig_w * scale), res_h), Image.LANCZOS) 717 + 718 + # Create comparison canvas 719 + gap = 20 720 + comp_w = orig_scaled.width + res_w + gap 721 + comp = Image.new("RGB", (comp_w, res_h + 40), (30, 30, 30)) 722 + 723 + # Paste images 724 + comp.paste(orig_scaled, (0, 0)) 725 + comp.paste(result, (orig_scaled.width + gap, 0)) 726 + 727 + # Add labels 728 + draw = ImageDraw.Draw(comp) 729 + draw.text((orig_scaled.width // 2 - 30, res_h + 5), "BEFORE", fill=(180, 180, 180)) 730 + draw.text((orig_scaled.width + gap + res_w // 2 - 20, res_h + 5), "AFTER", fill=(180, 180, 180)) 731 + 732 + comp.save(str(comparison_path), quality=95) 733 + print(f"Comparison saved to {comparison_path}") 734 + 735 + 736 + def main(): 737 + if len(sys.argv) < 2: 738 + print(f"Usage: {sys.argv[0]} <input.png> [output.png] [--debug] [--compare]") 739 + sys.exit(1) 740 + 741 + input_path = Path(sys.argv[1]) 742 + 743 + # Find output path 744 + positional_args = [a for a in sys.argv[2:] if not a.startswith("--")] 745 + output_path = Path(positional_args[0]) if positional_args else \ 746 + input_path.with_name(input_path.stem + "_photorealistic.png") 747 + 748 + debug = "--debug" in sys.argv 749 + compare = "--compare" in sys.argv 750 + 751 + if debug: 752 + save_debug_stages(input_path, output_path.parent / "debug") 753 + 754 + postprocess_photorealistic(input_path, output_path) 755 + 756 + if compare: 757 + comp_path = output_path.with_name(output_path.stem + "_comparison.png") 758 + create_comparison(input_path, output_path, comp_path) 759 + 760 + print("Done!") 761 + 762 + 763 + if __name__ == "__main__": 764 + main()
+235
worker/pipeline/svg2patch.py
··· 1 + #!/usr/bin/env python3 2 + """ 3 + svg2patch: Add inkstitch embroidery parameters to any SVG and render as a patch. 4 + 5 + Takes a clean SVG (from vtracer, manual design, etc.) and: 6 + 1. Adds inkstitch namespace + fill parameters to all paths 7 + 2. Sets document dimensions to mm 8 + 3. Renders with inkstitch for realistic stitch simulation 9 + 4. Optionally runs photorealistic post-processing 10 + """ 11 + 12 + import sys 13 + import os 14 + import subprocess 15 + import argparse 16 + from pathlib import Path 17 + from lxml import etree 18 + 19 + INKSTITCH_NS = "http://inkstitch.org/namespace" 20 + SVG_NS = "http://www.w3.org/2000/svg" 21 + INKSCAPE_NS = "http://www.inkscape.org/namespaces/inkscape" 22 + SODIPODI_NS = "http://sodipodi.sourceforge.net/DTD/sodipodi-0.0.dtd" 23 + 24 + INKSTITCH_BIN = os.environ.get( 25 + "INKSTITCH_BIN", 26 + os.path.expanduser( 27 + "~/Library/Application Support/org.inkscape.Inkscape" 28 + "/config/inkscape/extensions/inkstitch.app/Contents/MacOS/inkstitch" 29 + ), 30 + ) 31 + 32 + PATCH_WIDTH_MM = 80.0 33 + 34 + 35 + def add_inkstitch_params(svg_path, output_svg_path, border_color=None): 36 + """Add inkstitch embroidery parameters to all paths in an SVG.""" 37 + etree.register_namespace("inkstitch", INKSTITCH_NS) 38 + etree.register_namespace("inkscape", INKSCAPE_NS) 39 + etree.register_namespace("sodipodi", SODIPODI_NS) 40 + 41 + tree = etree.parse(str(svg_path)) 42 + root = tree.getroot() 43 + 44 + nsmap = dict(root.nsmap) 45 + nsmap["inkstitch"] = INKSTITCH_NS 46 + nsmap["inkscape"] = INKSCAPE_NS 47 + nsmap["sodipodi"] = SODIPODI_NS 48 + new_root = etree.Element(root.tag, nsmap=nsmap) 49 + new_root.attrib.update(root.attrib) 50 + new_root.text = root.text 51 + new_root.tail = root.tail 52 + for child in root: 53 + new_root.append(child) 54 + root = new_root 55 + 56 + # get original dimensions from viewBox or width/height 57 + viewbox = root.get("viewBox") 58 + if viewbox: 59 + parts = viewbox.split() 60 + vb_w = float(parts[2]) - float(parts[0]) 61 + vb_h = float(parts[3]) - float(parts[1]) 62 + else: 63 + vb_w = float(root.get("width", "100").replace("px", "").replace("mm", "")) 64 + vb_h = float(root.get("height", "100").replace("px", "").replace("mm", "")) 65 + 66 + # set dimensions to mm 67 + scale = PATCH_WIDTH_MM / vb_w 68 + width_mm = PATCH_WIDTH_MM 69 + height_mm = vb_h * scale 70 + root.set("width", f"{width_mm}mm") 71 + root.set("height", f"{height_mm}mm") 72 + 73 + # add namedview 74 + existing_nv = root.find("{%s}namedview" % SODIPODI_NS) 75 + if existing_nv is None: 76 + nv = etree.SubElement(root, "{%s}namedview" % SODIPODI_NS) 77 + nv.set("{%s}document-units" % INKSCAPE_NS, "mm") 78 + 79 + # add version metadata 80 + existing_meta = root.find("{%s}metadata" % SVG_NS) 81 + if existing_meta is None: 82 + existing_meta = root.find("metadata") 83 + if existing_meta is None: 84 + existing_meta = etree.SubElement(root, "metadata") 85 + version_el = existing_meta.find("{%s}inkstitch_svg_version" % INKSTITCH_NS) 86 + if version_el is None: 87 + version_el = etree.SubElement(existing_meta, "{%s}inkstitch_svg_version" % INKSTITCH_NS) 88 + version_el.text = "3" 89 + 90 + # find all paths/shapes and add inkstitch params 91 + all_elements = root.iter() 92 + shape_tags = { 93 + "{%s}path" % SVG_NS, "{%s}circle" % SVG_NS, "{%s}ellipse" % SVG_NS, 94 + "{%s}rect" % SVG_NS, "{%s}polygon" % SVG_NS, 95 + "path", "circle", "ellipse", "rect", "polygon", 96 + } 97 + 98 + element_count = 0 99 + for el in all_elements: 100 + tag = el.tag 101 + if tag not in shape_tags: 102 + continue 103 + 104 + # get fill color from style or fill attribute 105 + style = el.get("style", "") 106 + fill = el.get("fill", "") 107 + 108 + if "fill:none" in style or fill == "none": 109 + continue 110 + if "display:none" in style: 111 + continue 112 + 113 + # convert fill attribute to style if needed 114 + if fill and "fill:" not in style: 115 + if style: 116 + el.set("style", f"fill:{fill};stroke:none;{style}") 117 + else: 118 + el.set("style", f"fill:{fill};stroke:none") 119 + if el.get("fill"): 120 + del el.attrib["fill"] 121 + elif not fill and "fill:" not in style: 122 + continue 123 + 124 + # ensure stroke:none is in style 125 + current_style = el.get("style", "") 126 + if "stroke:" not in current_style: 127 + el.set("style", current_style.rstrip(";") + ";stroke:none") 128 + 129 + # add inkstitch fill parameters 130 + angle = (30 + element_count * 23) % 180 131 + el.set("{%s}fill_method" % INKSTITCH_NS, "auto_fill") 132 + el.set("{%s}fill_underlay" % INKSTITCH_NS, "true") 133 + el.set("{%s}fill_underlay_angle" % INKSTITCH_NS, str((angle + 90) % 360)) 134 + el.set("{%s}angle" % INKSTITCH_NS, str(angle)) 135 + el.set("{%s}row_spacing_mm" % INKSTITCH_NS, "0.25") 136 + el.set("{%s}max_stitch_length_mm" % INKSTITCH_NS, "3.0") 137 + el.set("{%s}staggers" % INKSTITCH_NS, "4") 138 + 139 + element_count += 1 140 + 141 + # add border if requested 142 + if border_color: 143 + if viewbox: 144 + parts = [float(p) for p in viewbox.split()] 145 + origin_x, origin_y = parts[0], parts[1] 146 + else: 147 + origin_x, origin_y = 0.0, 0.0 148 + 149 + pad = vb_w * 0.06 150 + border = etree.Element("rect") 151 + border.set("x", str(origin_x - pad)) 152 + border.set("y", str(origin_y - pad)) 153 + border.set("width", str(vb_w + pad * 2)) 154 + border.set("height", str(vb_h + pad * 2)) 155 + border.set("rx", str(pad * 0.8)) 156 + border.set("ry", str(pad * 0.8)) 157 + border.set("style", f"fill:{border_color};stroke:none") 158 + border.set("{%s}fill_method" % INKSTITCH_NS, "auto_fill") 159 + border.set("{%s}fill_underlay" % INKSTITCH_NS, "true") 160 + border.set("{%s}angle" % INKSTITCH_NS, "90") 161 + border.set("{%s}row_spacing_mm" % INKSTITCH_NS, "0.2") 162 + border.set("{%s}max_stitch_length_mm" % INKSTITCH_NS, "2.5") 163 + border.set("{%s}staggers" % INKSTITCH_NS, "4") 164 + root.insert(0, border) 165 + 166 + root.set("viewBox", f"{origin_x-pad} {origin_y-pad} {vb_w+pad*2} {vb_h+pad*2}") 167 + new_scale = PATCH_WIDTH_MM / (vb_w + pad * 2) 168 + root.set("width", f"{(vb_w + pad*2) * new_scale}mm") 169 + root.set("height", f"{(vb_h + pad*2) * new_scale}mm") 170 + element_count += 1 171 + 172 + etree.ElementTree(root).write(str(output_svg_path), xml_declaration=True, encoding="utf-8", pretty_print=True) 173 + print(f" {element_count} elements parameterized") 174 + return output_svg_path 175 + 176 + 177 + def render_inkstitch(svg_path, output_png): 178 + """Render with inkstitch realistic PNG.""" 179 + with open(output_png, "wb") as f: 180 + result = subprocess.run( 181 + [INKSTITCH_BIN, "--extension=png_realistic", str(svg_path)], 182 + stdout=f, stderr=subprocess.PIPE, timeout=300, 183 + ) 184 + if result.returncode != 0: 185 + print(f" inkstitch error: {result.stderr.decode()[:300]}", file=sys.stderr) 186 + if output_png.exists(): 187 + output_png.unlink() 188 + return False 189 + return output_png.exists() and output_png.stat().st_size > 0 190 + 191 + 192 + def main(): 193 + parser = argparse.ArgumentParser(description="Add inkstitch params to SVG and render as patch") 194 + parser.add_argument("input", help="Input SVG path") 195 + parser.add_argument("output", nargs="?", help="Output PNG path") 196 + parser.add_argument("-b", "--border-color", default="#0a0a14", help="Border color (default: #0a0a14)") 197 + parser.add_argument("--no-border", action="store_true", help="Skip border") 198 + parser.add_argument("--no-postprocess", action="store_true", help="Skip photorealistic post-processing") 199 + args = parser.parse_args() 200 + 201 + input_path = Path(args.input) 202 + output_path = Path(args.output) if args.output else input_path.with_name(input_path.stem + "_patch.png") 203 + 204 + border = None if args.no_border else args.border_color 205 + 206 + print(f"Adding inkstitch params to {input_path}...") 207 + embroidery_svg = output_path.with_suffix(".svg") 208 + add_inkstitch_params(input_path, embroidery_svg, border_color=border) 209 + 210 + print(f"Rendering with inkstitch...") 211 + if render_inkstitch(embroidery_svg, output_path): 212 + print(f" Stitch render: {output_path} ({output_path.stat().st_size // 1024}KB)") 213 + else: 214 + print(" inkstitch render failed") 215 + return 216 + 217 + if not args.no_postprocess: 218 + postprocess_script = Path(__file__).parent / "photorealistic.py" 219 + if postprocess_script.exists(): 220 + final_path = output_path.with_name(output_path.stem.replace("_patch", "") + "_final.png") 221 + print(f"Post-processing...") 222 + subprocess.run( 223 + [sys.executable, str(postprocess_script), str(output_path), str(final_path)], 224 + timeout=120, 225 + ) 226 + if final_path.exists(): 227 + print(f" Final: {final_path}") 228 + else: 229 + print(" (photorealistic.py not found, skipping post-processing)") 230 + 231 + print("Done!") 232 + 233 + 234 + if __name__ == "__main__": 235 + main()
+6
worker/requirements.txt
··· 1 + redis==5.0.0 2 + lxml==5.2.0 3 + numpy==1.26.0 4 + scipy==1.13.0 5 + Pillow==10.4.0 6 + vtracer==0.6.15
+76
worker/worker.py
··· 1 + import os 2 + import json 3 + import tempfile 4 + from pathlib import Path 5 + 6 + import redis 7 + 8 + from pipeline.convert import convert, convert_svg 9 + 10 + REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379") 11 + RESULT_TTL = int(os.environ.get("RESULT_TTL", "3600")) 12 + QUEUE_NAME = "patches" 13 + 14 + 15 + def main(): 16 + conn = redis.from_url(REDIS_URL) 17 + conn.ping() 18 + print(f"Worker ready, listening on queue '{QUEUE_NAME}'") 19 + 20 + while True: 21 + _, payload_bytes = conn.brpop(QUEUE_NAME) 22 + job = json.loads(payload_bytes) 23 + job_id = job["job_id"] 24 + print(f"Processing {job_id}") 25 + 26 + try: 27 + run_pipeline( 28 + conn, job_id, 29 + border_color=job.get("border_color", "#0a0a14"), 30 + color_precision=job.get("color_precision", 8), 31 + postprocess=job.get("postprocess", True), 32 + ) 33 + print(f" done") 34 + except Exception as e: 35 + print(f" failed: {e}") 36 + 37 + 38 + def run_pipeline(conn, job_id, border_color="#0a0a14", 39 + color_precision=8, postprocess=True): 40 + try: 41 + input_bytes = conn.get(f"job:{job_id}:input") 42 + ext = (conn.get(f"job:{job_id}:ext") or b"png").decode() 43 + if not input_bytes: 44 + raise RuntimeError("input not found in redis") 45 + 46 + is_svg = ext.lower() == "svg" 47 + 48 + with tempfile.TemporaryDirectory(prefix="p2p_") as tmpdir: 49 + tmpdir = Path(tmpdir) 50 + input_path = tmpdir / f"input.{ext}" 51 + input_path.write_bytes(input_bytes) 52 + output_path = tmpdir / "patch.png" 53 + 54 + if is_svg: 55 + convert_svg(str(input_path), str(output_path), 56 + border_color=border_color, postprocess=postprocess) 57 + else: 58 + convert(str(input_path), str(output_path), 59 + border_color=border_color, 60 + color_precision=color_precision, 61 + postprocess=postprocess) 62 + 63 + result_bytes = output_path.read_bytes() 64 + 65 + conn.setex(f"job:{job_id}:result", RESULT_TTL, result_bytes) 66 + conn.setex(f"job:{job_id}:status", RESULT_TTL, "complete") 67 + conn.delete(f"job:{job_id}:input") 68 + 69 + except Exception as e: 70 + conn.setex(f"job:{job_id}:status", RESULT_TTL, "failed") 71 + conn.setex(f"job:{job_id}:error", RESULT_TTL, str(e)) 72 + raise 73 + 74 + 75 + if __name__ == "__main__": 76 + main()