diff --git a/Dockerfile b/Dockerfile deleted file mode 100644 index da9039d..0000000 --- a/Dockerfile +++ /dev/null @@ -1,32 +0,0 @@ -FROM python:3.11-slim - -RUN apt-get update && apt-get install -y --no-install-recommends \ - wget xz-utils ca-certificates \ - libgtk-3-0t64 libdrm2 libgbm1 libxcb-dri3-0 \ - inkscape \ - && rm -rf /var/lib/apt/lists/* - -# install inkstitch (detect arch) -ARG INKSTITCH_VERSION=3.2.2 -RUN ARCH=$(dpkg --print-architecture) && \ - if [ "$ARCH" = "arm64" ]; then INKARCH="aarch64"; else INKARCH="x86_64"; fi && \ - wget -q "https://github.com/inkstitch/inkstitch/releases/download/v${INKSTITCH_VERSION}/inkstitch-${INKSTITCH_VERSION}-linux-${INKARCH}.tar.xz" \ - -O /tmp/inkstitch.tar.xz && \ - mkdir -p /opt/inkstitch && \ - tar xJf /tmp/inkstitch.tar.xz -C /opt/inkstitch && \ - rm /tmp/inkstitch.tar.xz -ENV INKSTITCH_BIN=/opt/inkstitch/inkstitch/bin/inkstitch - -WORKDIR /app -COPY requirements.txt . -RUN pip install --no-cache-dir -r requirements.txt vtracer - -COPY pipeline/ pipeline/ -COPY server.py worker.py ./ - -ENV RESULTS_DIR=/data/results -ENV VTRACER_BIN=vtracer - -ENV MODE=api -EXPOSE 8000 -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' diff --git a/Dockerfile.web b/Dockerfile.web index ea27e9b..35faf0b 100644 --- a/Dockerfile.web +++ b/Dockerfile.web @@ -1,16 +1,12 @@ -FROM python:3.11-slim +FROM golang:1.24-alpine AS build -WORKDIR /app +WORKDIR /src +COPY web/go.mod web/go.sum ./ +RUN go mod download +COPY web/main.go . +RUN CGO_ENABLED=0 go build -o /server . -RUN pip install --no-cache-dir \ - fastapi==0.115.0 \ - uvicorn==0.30.0 \ - python-multipart==0.0.9 \ - redis==5.0.0 \ - rq==1.16.0 - -COPY server.py . - -ENV RESULTS_DIR=/data/results +FROM alpine:3.20 +COPY --from=build /server /server EXPOSE 8000 -CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8000"] +CMD ["/server"] diff --git a/Dockerfile.worker b/Dockerfile.worker index e652415..99a6b6d 100644 --- a/Dockerfile.worker +++ b/Dockerfile.worker @@ -1,9 +1,8 @@ -FROM python:3.11-slim +# Stage 1: download and strip inkstitch +FROM python:3.11-slim AS builder RUN apt-get update && apt-get install -y --no-install-recommends \ wget xz-utils ca-certificates \ - libgtk-3-0t64 libdrm2 libgbm1 libxcb-dri3-0 \ - inkscape \ && rm -rf /var/lib/apt/lists/* ARG INKSTITCH_VERSION=3.2.2 @@ -13,15 +12,31 @@ RUN ARCH=$(dpkg --print-architecture) && \ -O /tmp/inkstitch.tar.xz && \ mkdir -p /opt/inkstitch && \ tar xJf /tmp/inkstitch.tar.xz -C /opt/inkstitch && \ - rm /tmp/inkstitch.tar.xz + rm /tmp/inkstitch.tar.xz && \ + rm -rf /opt/inkstitch/inkstitch/fonts \ + /opt/inkstitch/inkstitch/bin/locales \ + /opt/inkstitch/inkstitch/inx \ + /opt/inkstitch/inkstitch/palettes \ + /opt/inkstitch/inkstitch/tiles \ + /opt/inkstitch/inkstitch/addons \ + /opt/inkstitch/inkstitch/symbols \ + /opt/inkstitch/inkstitch/dbus + +# Stage 2: runtime +FROM python:3.11-slim + +RUN apt-get update && apt-get install -y --no-install-recommends \ + inkscape libwayland-cursor0 libdrm2 libgbm1 libxcb-dri3-0 \ + && rm -rf /var/lib/apt/lists/* + +COPY --from=builder /opt/inkstitch /opt/inkstitch ENV INKSTITCH_BIN=/opt/inkstitch/inkstitch/bin/inkstitch WORKDIR /app -COPY requirements.txt . +COPY worker/requirements.txt . RUN pip install --no-cache-dir -r requirements.txt -COPY pipeline/ pipeline/ -COPY worker.py . +COPY worker/pipeline/ pipeline/ +COPY worker/worker.py . -ENV RESULTS_DIR=/data/results -CMD ["rq", "worker", "patches", "--url", "redis://redis:6379"] +CMD ["python", "worker.py"] diff --git a/docker-compose.yml b/docker-compose.yml index 7dd5099..d144c4d 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -5,27 +5,21 @@ services: - "6379:6379" api: - build: . + build: + context: . + dockerfile: Dockerfile.web ports: - "8000:8000" environment: - REDIS_URL=redis://redis:6379 - - RESULTS_DIR=/data/results - volumes: - - results:/data/results depends_on: - redis worker: - build: . - command: rq worker patches --url redis://redis:6379 + build: + context: . + dockerfile: Dockerfile.worker environment: - REDIS_URL=redis://redis:6379 - - RESULTS_DIR=/data/results - volumes: - - results:/data/results depends_on: - redis - -volumes: - results: diff --git a/server.py b/server.py deleted file mode 100644 index 1504ec2..0000000 --- a/server.py +++ /dev/null @@ -1,88 +0,0 @@ -""" -pic-to-patch API server. - -POST /patch — upload image, returns {job_id} -GET /jobs/{job_id} — returns {status, result_url?} -GET /jobs/{job_id}/result — returns the PNG directly - -Jobs are queued via Redis (rq). Results are stored in Redis as blobs. -""" - -import os -import uuid - -from fastapi import FastAPI, UploadFile, HTTPException -from fastapi.responses import Response -import redis -import rq - -REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379") -RESULT_TTL = int(os.environ.get("RESULT_TTL", "3600")) - -app = FastAPI(title="pic-to-patch", version="0.2.0") -redis_conn = redis.from_url(REDIS_URL) -queue = rq.Queue("patches", connection=redis_conn) - - -@app.post("/patch") -async def create_patch( - file: UploadFile, - border_color: str = "#0a0a14", - color_precision: int = 8, - postprocess: bool = True, -): - job_id = str(uuid.uuid4()) - input_bytes = await file.read() - ext = (file.filename or "input.png").rsplit(".", 1)[-1] or "png" - - redis_conn.setex(f"job:{job_id}:input", RESULT_TTL, input_bytes) - redis_conn.setex(f"job:{job_id}:ext", RESULT_TTL, ext) - redis_conn.setex(f"job:{job_id}:status", RESULT_TTL, "processing") - - queue.enqueue( - "worker.run_pipeline", - job_id=job_id, - border_color=border_color, - color_precision=color_precision, - postprocess=postprocess, - job_timeout=300, - result_ttl=RESULT_TTL, - ) - - return {"job_id": job_id} - - -@app.get("/jobs/{job_id}") -async def get_job(job_id: str): - status = redis_conn.get(f"job:{job_id}:status") - if status is None: - raise HTTPException(404, "job not found") - - status = status.decode() - resp = {"status": status} - - if status == "complete": - resp["result_url"] = f"/jobs/{job_id}/result" - elif status == "failed": - error = redis_conn.get(f"job:{job_id}:error") - resp["error"] = error.decode() if error else "unknown error" - - return resp - - -@app.get("/jobs/{job_id}/result") -async def get_result(job_id: str): - data = redis_conn.get(f"job:{job_id}:result") - if data is None: - raise HTTPException(404, "result not ready") - return Response(content=data, media_type="image/png", - headers={"Content-Disposition": f'inline; filename="{job_id}.png"'}) - - -@app.get("/health") -async def health(): - try: - redis_conn.ping() - except Exception: - raise HTTPException(503, "redis unavailable") - return {"status": "ok"} diff --git a/web/go.mod b/web/go.mod new file mode 100644 index 0000000..d1b9c2d --- /dev/null +++ b/web/go.mod @@ -0,0 +1,13 @@ +module pic-to-patch-web + +go 1.24 + +require ( + github.com/google/uuid v1.6.0 + github.com/redis/go-redis/v9 v9.19.0 +) + +require ( + github.com/cespare/xxhash/v2 v2.3.0 // indirect + go.uber.org/atomic v1.11.0 // indirect +) diff --git a/web/go.sum b/web/go.sum new file mode 100644 index 0000000..a7c978e --- /dev/null +++ b/web/go.sum @@ -0,0 +1,24 @@ +github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= +github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= +github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= +github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE= +github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/redis/go-redis/v9 v9.19.0 h1:XPVaaPSnG6RhYf7p+rmSa9zZfeVAnWsH5h3lxthOm/k= +github.com/redis/go-redis/v9 v9.19.0/go.mod h1:v/M13XI1PVCDcm01VtPFOADfZtHf8YW3baQf57KlIkA= +github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= +github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= +go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= +go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= +golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc= +golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= diff --git a/web/main.go b/web/main.go new file mode 100644 index 0000000..149c73f --- /dev/null +++ b/web/main.go @@ -0,0 +1,184 @@ +package main + +import ( + "context" + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "os" + "strconv" + "strings" + "time" + + "github.com/google/uuid" + "github.com/redis/go-redis/v9" +) + +var ( + rdb *redis.Client + resultTTL time.Duration +) + +type QueuePayload struct { + JobID string `json:"job_id"` + BorderColor string `json:"border_color"` + ColorPrecision int `json:"color_precision"` + Postprocess bool `json:"postprocess"` +} + +func main() { + redisURL := env("REDIS_URL", "redis://localhost:6379") + ttlSec, _ := strconv.Atoi(env("RESULT_TTL", "3600")) + resultTTL = time.Duration(ttlSec) * time.Second + + opt, err := redis.ParseURL(redisURL) + if err != nil { + log.Fatalf("bad REDIS_URL: %v", err) + } + rdb = redis.NewClient(opt) + + mux := http.NewServeMux() + mux.HandleFunc("POST /patch", handleCreatePatch) + mux.HandleFunc("GET /jobs/{job_id}", handleGetJob) + mux.HandleFunc("GET /jobs/{job_id}/result", handleGetResult) + mux.HandleFunc("GET /health", handleHealth) + + log.Println("listening on :8000") + log.Fatal(http.ListenAndServe(":8000", mux)) +} + +func handleCreatePatch(w http.ResponseWriter, r *http.Request) { + file, header, err := r.FormFile("file") + if err != nil { + http.Error(w, "missing file", http.StatusBadRequest) + return + } + defer file.Close() + + inputBytes, err := io.ReadAll(file) + if err != nil { + http.Error(w, "failed to read file", http.StatusInternalServerError) + return + } + + filename := "input.png" + if header.Filename != "" { + filename = header.Filename + } + ext := "png" + if i := strings.LastIndex(filename, "."); i >= 0 && i+1 < len(filename) { + ext = filename[i+1:] + } + + borderColor := r.FormValue("border_color") + if borderColor == "" { + borderColor = "#0a0a14" + } + + colorPrecision := 8 + if v := r.FormValue("color_precision"); v != "" { + colorPrecision, _ = strconv.Atoi(v) + } + + postprocess := true + if v := r.FormValue("postprocess"); v != "" { + postprocess = v != "false" && v != "0" && v != "no" + } + + jobID := uuid.New().String() + ctx := context.Background() + + pipe := rdb.Pipeline() + pipe.SetEx(ctx, fmt.Sprintf("job:%s:input", jobID), inputBytes, resultTTL) + pipe.SetEx(ctx, fmt.Sprintf("job:%s:ext", jobID), ext, resultTTL) + pipe.SetEx(ctx, fmt.Sprintf("job:%s:status", jobID), "processing", resultTTL) + if _, err := pipe.Exec(ctx); err != nil { + http.Error(w, "redis error", http.StatusInternalServerError) + return + } + + payload := QueuePayload{ + JobID: jobID, + BorderColor: borderColor, + ColorPrecision: colorPrecision, + Postprocess: postprocess, + } + payloadJSON, _ := json.Marshal(payload) + + if err := rdb.LPush(ctx, "patches", payloadJSON).Err(); err != nil { + http.Error(w, "queue error", http.StatusInternalServerError) + return + } + + writeJSON(w, http.StatusOK, map[string]string{"job_id": jobID}) +} + +func handleGetJob(w http.ResponseWriter, r *http.Request) { + jobID := r.PathValue("job_id") + ctx := context.Background() + + status, err := rdb.Get(ctx, fmt.Sprintf("job:%s:status", jobID)).Result() + if err == redis.Nil { + http.Error(w, `{"detail":"job not found"}`, http.StatusNotFound) + return + } else if err != nil { + http.Error(w, "redis error", http.StatusInternalServerError) + return + } + + resp := map[string]interface{}{"status": status} + + if status == "complete" { + resp["result_url"] = fmt.Sprintf("/jobs/%s/result", jobID) + } else if status == "failed" { + errMsg, _ := rdb.Get(ctx, fmt.Sprintf("job:%s:error", jobID)).Result() + if errMsg == "" { + errMsg = "unknown error" + } + resp["error"] = errMsg + } + + writeJSON(w, http.StatusOK, resp) +} + +func handleGetResult(w http.ResponseWriter, r *http.Request) { + jobID := r.PathValue("job_id") + ctx := context.Background() + + data, err := rdb.Get(ctx, fmt.Sprintf("job:%s:result", jobID)).Bytes() + if err == redis.Nil { + http.Error(w, `{"detail":"result not ready"}`, http.StatusNotFound) + return + } else if err != nil { + http.Error(w, "redis error", http.StatusInternalServerError) + return + } + + w.Header().Set("Content-Type", "image/png") + w.Header().Set("Content-Disposition", fmt.Sprintf(`inline; filename="%s.png"`, jobID)) + w.Write(data) +} + +func handleHealth(w http.ResponseWriter, r *http.Request) { + ctx := context.Background() + if err := rdb.Ping(ctx).Err(); err != nil { + http.Error(w, `{"detail":"redis unavailable"}`, http.StatusServiceUnavailable) + return + } + writeJSON(w, http.StatusOK, map[string]string{"status": "ok"}) +} + +func writeJSON(w http.ResponseWriter, status int, v interface{}) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + json.NewEncoder(w).Encode(v) +} + +func env(key, fallback string) string { + if v := os.Getenv(key); v != "" { + return v + } + return fallback +} diff --git a/worker.py b/worker.py deleted file mode 100644 index f397151..0000000 --- a/worker.py +++ /dev/null @@ -1,53 +0,0 @@ -""" -pic-to-patch worker. Picks jobs from Redis queue and runs the pipeline. - -Run with: rq worker patches --url redis://localhost:6379 -""" - -import os -import tempfile -from pathlib import Path - -import redis - -from pipeline.convert import convert, convert_svg - -REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379") -RESULT_TTL = int(os.environ.get("RESULT_TTL", "3600")) -redis_conn = redis.from_url(REDIS_URL) - - -def run_pipeline(job_id, border_color="#0a0a14", - color_precision=8, postprocess=True): - try: - input_bytes = redis_conn.get(f"job:{job_id}:input") - ext = (redis_conn.get(f"job:{job_id}:ext") or b"png").decode() - if not input_bytes: - raise RuntimeError("input not found in redis") - - is_svg = ext.lower() == "svg" - - with tempfile.TemporaryDirectory(prefix="p2p_") as tmpdir: - tmpdir = Path(tmpdir) - input_path = tmpdir / f"input.{ext}" - input_path.write_bytes(input_bytes) - output_path = tmpdir / "patch.png" - - if is_svg: - convert_svg(str(input_path), str(output_path), - border_color=border_color, postprocess=postprocess) - else: - convert(str(input_path), str(output_path), - border_color=border_color, color_precision=color_precision, - postprocess=postprocess) - - result_bytes = output_path.read_bytes() - - redis_conn.setex(f"job:{job_id}:result", RESULT_TTL, result_bytes) - redis_conn.setex(f"job:{job_id}:status", RESULT_TTL, "complete") - redis_conn.delete(f"job:{job_id}:input") - - except Exception as e: - redis_conn.setex(f"job:{job_id}:status", RESULT_TTL, "failed") - redis_conn.setex(f"job:{job_id}:error", RESULT_TTL, str(e)) - raise diff --git a/worker/pipeline/__init__.py b/worker/pipeline/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/worker/pipeline/convert.py b/worker/pipeline/convert.py new file mode 100644 index 0000000..f5b3158 --- /dev/null +++ b/worker/pipeline/convert.py @@ -0,0 +1,116 @@ +""" +Core pipeline: image → embroidered patch PNG. + +This is the throwaway python version. The interface is: + convert(input_path, output_path, border_color="#0a0a14", color_precision=8) + +That's it. Everything else is internal. +""" + +import os +import subprocess +import tempfile +from pathlib import Path + +from PIL import Image +import vtracer + +from .svg2patch import add_inkstitch_params, render_inkstitch + +MAX_IMAGE_DIM = int(os.environ.get("MAX_IMAGE_DIM", "500")) + + +def convert(input_path, output_path, border_color="#0a0a14", color_precision=8, postprocess=True): + input_path = Path(input_path) + output_path = Path(output_path) + + with tempfile.TemporaryDirectory(prefix="p2p_") as tmpdir: + tmpdir = Path(tmpdir) + + resized = tmpdir / "resized.png" + _resize(input_path, resized) + + vtracer_svg = tmpdir / "vectorized.svg" + _vectorize(resized, vtracer_svg, color_precision) + + embroidery_svg = tmpdir / "embroidery.svg" + add_inkstitch_params(vtracer_svg, embroidery_svg, border_color=border_color) + + stitch_png = tmpdir / "stitch.png" + if not render_inkstitch(embroidery_svg, stitch_png): + raise RuntimeError("inkstitch render failed") + + if postprocess: + _postprocess(stitch_png, output_path) + else: + import shutil + shutil.copy2(stitch_png, output_path) + + return output_path + + +def convert_svg(input_svg, output_path, border_color="#0a0a14", postprocess=True): + input_svg = Path(input_svg) + output_path = Path(output_path) + + with tempfile.TemporaryDirectory(prefix="p2p_") as tmpdir: + tmpdir = Path(tmpdir) + + embroidery_svg = tmpdir / "embroidery.svg" + add_inkstitch_params(input_svg, embroidery_svg, border_color=border_color) + + stitch_png = tmpdir / "stitch.png" + if not render_inkstitch(embroidery_svg, stitch_png): + raise RuntimeError("inkstitch render failed") + + if postprocess: + _postprocess(stitch_png, output_path) + else: + import shutil + shutil.copy2(stitch_png, output_path) + + return output_path + + +def _resize(input_path, output_path): + img = Image.open(input_path) + if img.mode == "P": + img = img.convert("RGBA") + w, h = img.size + if max(w, h) > MAX_IMAGE_DIM: + ratio = MAX_IMAGE_DIM / max(w, h) + img = img.resize((int(w * ratio), int(h * ratio)), Image.Resampling.LANCZOS) + img.save(str(output_path)) + + +def _vectorize(input_image, output_svg, color_precision=8): + vtracer.convert_image_to_svg_py( + image_path=str(input_image), + out_path=str(output_svg), + colormode="color", + hierarchical="stacked", + filter_speckle=4, + color_precision=color_precision, + corner_threshold=60, + length_threshold=4, + splice_threshold=45, + ) + if not output_svg.exists(): + raise RuntimeError("vtracer produced no output") + + +def _postprocess(stitch_png, output_path): + script = Path(__file__).parent / "photorealistic.py" + if not script.exists(): + import shutil + shutil.copy2(stitch_png, output_path) + return + + import sys + result = subprocess.run( + [sys.executable, str(script), str(stitch_png), str(output_path)], + capture_output=True, text=True, timeout=120, + ) + if result.returncode != 0 or not output_path.exists(): + import shutil + shutil.copy2(stitch_png, output_path) diff --git a/worker/pipeline/photorealistic.py b/worker/pipeline/photorealistic.py new file mode 100644 index 0000000..ab92a19 --- /dev/null +++ b/worker/pipeline/photorealistic.py @@ -0,0 +1,764 @@ +#!/usr/bin/env python3 +""" +Advanced post-processing pipeline that makes inkstitch renders look like +photographs of real embroidered patches. + +Pipeline: + 1. Extract patch mask from white background + 2. Compute per-pixel normal map from stitch texture gradients + 3. Estimate local thread direction via structure tensor + 4. Apply Blinn-Phong shading with Kajiya-Kay anisotropic specular + 5. Add per-thread micro-highlight variation + 6. Generate felt/twill backing fabric texture + 7. Create realistic drop shadow (contact + cast) + 8. Simulate merrow/overlock edge border + 9. Add patch thickness bevel (embossed raised edge) + 10. Composite everything + 11. Photographic finishing: DOF, color grade, vignette, grain +""" + +import sys +import numpy as np +from pathlib import Path +from PIL import Image, ImageFilter, ImageEnhance, ImageDraw +from scipy import ndimage +from scipy.ndimage import gaussian_filter, uniform_filter + + +# ───────────────────────────────────────────── +# 1. Mask extraction +# ───────────────────────────────────────────── + +def extract_patch_mask(img_arr, threshold=235): + """ + Extract a mask of the embroidered region vs white background. + Returns float mask [0, 1] with anti-aliased soft edges. + """ + is_bg = np.all(img_arr[:, :, :3] > threshold, axis=2) + mask = (~is_bg).astype(np.float32) + + # Close small holes inside the patch + mask = ndimage.binary_closing(mask > 0.5, iterations=3).astype(np.float32) + # Dilate slightly to catch anti-aliased edge pixels + mask = ndimage.binary_dilation(mask > 0.5, iterations=1).astype(np.float32) + + # Soft edge via gaussian blur + mask_soft = gaussian_filter(mask, sigma=1.2) + return np.clip(mask_soft, 0, 1) + + +# ───────────────────────────────────────────── +# 2. Normal map computation +# ───────────────────────────────────────────── + +def compute_normal_map(img_arr, mask, strength=2.0): + """ + Derive per-pixel surface normals from stitch texture luminance gradients. + Multi-scale: fine captures thread ridges, medium captures stitch rows. + """ + lum = (0.299 * img_arr[:, :, 0] + + 0.587 * img_arr[:, :, 1] + + 0.114 * img_arr[:, :, 2]).astype(np.float64) + + # Fine scale - individual thread ridges + gx_fine = ndimage.sobel(lum, axis=1) + gy_fine = ndimage.sobel(lum, axis=0) + + # Medium scale - stitch row structure + lum_med = gaussian_filter(lum, sigma=1.2) + gx_med = ndimage.sobel(lum_med, axis=1) + gy_med = ndimage.sobel(lum_med, axis=0) + + # Coarse scale - broad curvature + lum_coarse = gaussian_filter(lum, sigma=3.0) + gx_coarse = ndimage.sobel(lum_coarse, axis=1) + gy_coarse = ndimage.sobel(lum_coarse, axis=0) + + gx = 0.5 * gx_fine + 0.35 * gx_med + 0.15 * gx_coarse + gy = 0.5 * gy_fine + 0.35 * gy_med + 0.15 * gy_coarse + + gx *= strength + gy *= strength + + h, w = lum.shape + normals = np.zeros((h, w, 3), dtype=np.float64) + normals[:, :, 0] = -gx + normals[:, :, 1] = -gy + normals[:, :, 2] = 255.0 + + length = np.sqrt(np.sum(normals ** 2, axis=2, keepdims=True)) + normals /= np.maximum(length, 1e-8) + normals *= mask[:, :, np.newaxis] + + return normals + + +# ───────────────────────────────────────────── +# 3. Thread direction estimation +# ───────────────────────────────────────────── + +def estimate_thread_direction(img_arr, mask, block_size=12): + """ + Estimate local thread direction via structure tensor analysis. + Returns angle map in radians for anisotropic specular. + """ + lum = (0.299 * img_arr[:, :, 0] + + 0.587 * img_arr[:, :, 1] + + 0.114 * img_arr[:, :, 2]).astype(np.float64) + + gx = ndimage.sobel(lum, axis=1) + gy = ndimage.sobel(lum, axis=0) + + sigma = block_size / 2.0 + Jxx = gaussian_filter(gx * gx, sigma=sigma) + Jxy = gaussian_filter(gx * gy, sigma=sigma) + Jyy = gaussian_filter(gy * gy, sigma=sigma) + + # Dominant gradient direction + angle = 0.5 * np.arctan2(2.0 * Jxy, Jxx - Jyy + 1e-10) + # Thread runs perpendicular to gradient + thread_angle = angle + np.pi / 2.0 + + # Also compute anisotropy strength (how directional the texture is) + trace = Jxx + Jyy + 1e-10 + diff = np.sqrt((Jxx - Jyy) ** 2 + 4 * Jxy ** 2) + anisotropy = diff / trace # 0 = isotropic, 1 = perfectly directional + + return thread_angle, anisotropy + + +# ───────────────────────────────────────────── +# 4. Blinn-Phong + anisotropic shading +# ───────────────────────────────────────────── + +def apply_lighting(img_arr, normals, mask, thread_angle, anisotropy, + light_dir=(0.25, -0.35, 0.90), + ambient=0.55, diffuse_strength=0.35, + specular_strength=0.12, shininess=18.0, + aniso_strength=0.08, aniso_shininess=30.0): + """ + Apply Blinn-Phong shading with Kajiya-Kay anisotropic specular. + + The key insight: we DON'T want to darken the patch too much. Real embroidery + is bright and saturated under normal lighting. The shading should add + subtle variation and highlights, not dramatically change overall brightness. + + We use a "detail lighting" approach: + final = original * (ambient + diffuse) + specular_white + where ambient+diffuse averages close to 1.0 across the image. + """ + result = img_arr.astype(np.float64).copy() + + light = np.array(light_dir, dtype=np.float64) + light /= np.linalg.norm(light) + + view = np.array([0.0, 0.0, 1.0]) + half_vec = light + view + half_vec /= np.linalg.norm(half_vec) + + # --- Diffuse --- + n_dot_l = (normals[:, :, 0] * light[0] + + normals[:, :, 1] * light[1] + + normals[:, :, 2] * light[2]) + n_dot_l = np.clip(n_dot_l, 0, 1) + + # Wrap diffuse to soften shadows (half-lambert) + diffuse = diffuse_strength * (n_dot_l * 0.5 + 0.5) + + # --- Isotropic specular --- + n_dot_h = (normals[:, :, 0] * half_vec[0] + + normals[:, :, 1] * half_vec[1] + + normals[:, :, 2] * half_vec[2]) + n_dot_h = np.clip(n_dot_h, 0, 1) + specular_iso = specular_strength * np.power(n_dot_h, shininess) + + # --- Anisotropic specular (Kajiya-Kay) --- + tangent_x = np.cos(thread_angle) + tangent_y = np.sin(thread_angle) + + t_dot_h = tangent_x * half_vec[0] + tangent_y * half_vec[1] + sin_th = np.sqrt(np.clip(1.0 - t_dot_h ** 2, 0, 1)) + # Modulate by local anisotropy - only apply aniso spec where texture is directional + specular_aniso = aniso_strength * np.power(sin_th, aniso_shininess) * anisotropy + + # --- Secondary broad specular (satin sheen on lighter regions) --- + # Lighter areas (white text, bright embroidery) have more specular + lum = np.mean(result[:, :, :3], axis=2) / 255.0 + brightness_boost = np.clip(lum - 0.4, 0, 0.6) / 0.6 # ramp from 0.4 to 1.0 + spec_broad = 0.06 * np.power(n_dot_h, 8.0) * brightness_boost + + # --- Combine --- + # Multiplicative lighting (affects color) + color_light = ambient + diffuse + # Ensure average is near 1.0 to preserve original brightness + color_light = np.clip(color_light, 0.35, 1.5) + + # Additive specular (white highlight) + spec_total = (specular_iso + specular_aniso + spec_broad) * mask + + # Apply + for c in range(3): + channel = result[:, :, c] + lit = channel * color_light + spec_total * 220.0 + result[:, :, c] = channel * (1.0 - mask) + lit * mask + + return np.clip(result, 0, 255) + + +# ───────────────────────────────────────────── +# 5. Per-thread micro-highlights +# ───────────────────────────────────────────── + +def compute_ambient_occlusion(img_arr, mask, radius=2.0, strength=0.15): + """ + Approximate screen-space ambient occlusion from stitch texture. + Stitch valleys (between thread rows) are slightly darker. + Computed by comparing local brightness to neighborhood average. + """ + lum = np.mean(img_arr[:, :, :3], axis=2) + local_avg = gaussian_filter(lum, sigma=radius) + ao = np.clip((local_avg - lum) / (local_avg + 1e-8), 0, 1) * strength + ao *= mask + return ao + + +def add_thread_microhighlights(img_arr, normals, mask, thread_angle, intensity=0.05): + """ + Individual threads catch light slightly differently. + Creates fine-grained brightness variation aligned with thread direction. + Also adds subtle shimmer for thread luster. + """ + h, w = img_arr.shape[:2] + noise = np.random.normal(0, 1, (h, w)) + + # Directional blur kernels + noise_along = gaussian_filter(noise, sigma=[0.3, 3.0]) + noise_across = gaussian_filter(noise, sigma=[3.0, 0.3]) + + # Blend using thread angle + cos_a = np.abs(np.cos(thread_angle)) + sin_a = np.abs(np.sin(thread_angle)) + total = cos_a + sin_a + 1e-8 + directional = (noise_along * cos_a + noise_across * sin_a) / total + + highlight = directional * intensity * mask + result = img_arr.copy() + for c in range(3): + result[:, :, c] *= (1.0 + highlight) + + # Fine per-pixel shimmer for thread luster + shimmer = np.random.normal(0, 0.015, (h, w)) + shimmer = gaussian_filter(shimmer, sigma=0.3) + for c in range(3): + result[:, :, c] *= (1.0 + shimmer * mask) + + return np.clip(result, 0, 255) + + +# ───────────────────────────────────────────── +# 6. Fabric backing texture +# ───────────────────────────────────────────── + +def generate_fabric_texture(width, height, color=(35, 35, 40), weave_scale=3): + """ + Generate realistic felt/twill fabric backing. + Multi-octave noise with diagonal weave pattern. + """ + arr = np.full((height, width, 3), color, dtype=np.float64) + + y_idx = np.arange(height)[:, None] + x_idx = np.arange(width)[None, :] + + # Twill diagonal weave + twill1 = np.sin(2 * np.pi * (x_idx + y_idx) / weave_scale) * 0.035 + twill2 = np.sin(2 * np.pi * (x_idx - y_idx) / (weave_scale * 1.3)) * 0.02 + twill3 = np.sin(2 * np.pi * (x_idx * 0.7 + y_idx * 1.3) / (weave_scale * 2)) * 0.015 + arr *= (1.0 + twill1 + twill2 + twill3)[:, :, np.newaxis] + + # Fine fiber noise + noise_fine = np.random.normal(0, 2.0, (height, width, 3)) + # Medium texture clumps + noise_med_small = np.random.normal(0, 1.2, (height // 2 + 1, width // 2 + 1, 3)) + noise_med = np.repeat(np.repeat(noise_med_small, 2, axis=0), 2, axis=1)[:height, :width, :] + noise_med = gaussian_filter(noise_med, sigma=1.0) + # Coarse color drift + noise_coarse_small = np.random.normal(0, 0.8, (height // 6 + 1, width // 6 + 1, 3)) + noise_coarse = np.repeat(np.repeat(noise_coarse_small, 6, axis=0), 6, axis=1)[:height, :width, :] + noise_coarse = gaussian_filter(noise_coarse, sigma=3.0) + + arr += noise_fine + noise_med + noise_coarse + + # Subtle large-scale brightness variation (fabric isn't perfectly uniform) + var_small = np.random.normal(0, 0.015, (height // 12 + 1, width // 12 + 1)) + variation = np.repeat(np.repeat(var_small, 12, axis=0), 12, axis=1)[:height, :width] + variation = gaussian_filter(variation, sigma=6.0) + arr *= (1.0 + variation[:, :, np.newaxis]) + + return np.clip(arr, 0, 255).astype(np.uint8) + + +# ───────────────────────────────────────────── +# 7. Shadow +# ───────────────────────────────────────────── + +def create_patch_shadow(mask, offset=(5, 6), blur_radius=10, opacity=0.55): + """ + Realistic two-layer shadow: + - Contact shadow: tight, dark, right at the edge + - Cast shadow: offset, soft, diffuse + """ + h, w = mask.shape + + # Cast shadow: shift the mask by offset and blur + shifted = ndimage.shift(mask, (offset[1], offset[0]), order=1, mode='constant', cval=0) + cast = gaussian_filter(shifted, sigma=blur_radius) * opacity + + # Contact shadow: unshifted tight edge glow + # Expand mask slightly, subtract original, blur tightly + dilated = gaussian_filter(mask, sigma=2.0) + contact_ring = np.clip(dilated - mask * 0.9, 0, 1) + contact = gaussian_filter(contact_ring, sigma=2.5) * 0.4 + + # Combine + shadow = np.maximum(cast, contact) + # Don't darken the patch interior + shadow *= np.clip(1.0 - mask, 0, 1) + + return np.clip(shadow, 0, 1) + + +# ───────────────────────────────────────────── +# 8. Merrow edge +# ───────────────────────────────────────────── + +def create_merrow_edge(mask, thickness=3): + """ + Simulate overlock/merrow stitch border around patch edge. + Returns (edge_mask, edge_normals) for 3D stitched border look. + """ + h, w = mask.shape + + # Get edge band via morphological gradient + hard_mask = (mask > 0.5).astype(np.float32) + dilated = ndimage.binary_dilation(hard_mask > 0.5, iterations=thickness).astype(np.float32) + eroded = ndimage.binary_erosion(hard_mask > 0.5, iterations=max(1, thickness // 2)).astype(np.float32) + edge_band = np.clip(dilated - eroded, 0, 1) + + # Anti-alias the edge band + edge_band = gaussian_filter(edge_band, sigma=0.6) + + # Add stitch texture pattern along the edge + y_idx = np.arange(h)[:, None].astype(np.float64) + x_idx = np.arange(w)[None, :].astype(np.float64) + + # Distance from center for radial stitch direction + cy, cx = h / 2.0, w / 2.0 + dy = y_idx - cy + dx = x_idx - cx + angle = np.arctan2(dy, dx) + + # Stitch pattern follows the edge circumferentially + # Use angle to create perpendicular stitches + dist = np.sqrt(dx ** 2 + dy ** 2) + stitch_freq = dist * 0.15 # scale frequency by radius + stitch_pattern = np.sin(stitch_freq + angle * 25) * 0.12 + 0.88 + stitch_pattern2 = np.cos(stitch_freq * 1.7 + angle * 18) * 0.06 + 0.94 + + edge_textured = edge_band * stitch_pattern * stitch_pattern2 + + return edge_textured, edge_band + + +# ───────────────────────────────────────────── +# 9. Patch thickness bevel +# ───────────────────────────────────────────── + +def create_edge_bevel(mask, bevel_width=6, light_dir=(0.3, -0.4)): + """ + Create a bevel/emboss effect at the patch edge to simulate thickness. + The patch is raised ~1-2mm above the fabric, creating a lit top edge + and shadowed bottom edge. + """ + # Compute distance from edge (inward) + hard_mask = (mask > 0.5).astype(np.float32) + dist = ndimage.distance_transform_edt(hard_mask) + dist_outside = ndimage.distance_transform_edt(1 - hard_mask) + + # Bevel height profile: ramps up at edge, flat in interior + bevel_height = np.clip(dist / bevel_width, 0, 1) + # Also slight ramp outside for the outer edge + bevel_height_out = np.clip(1.0 - dist_outside / (bevel_width * 0.5), 0, 1) * (1 - hard_mask) + + height = bevel_height + bevel_height_out + + # Compute lighting from height map + gx = ndimage.sobel(height, axis=1) + gy = ndimage.sobel(height, axis=0) + + # Directional lighting + lx, ly = light_dir + bevel_light = -(gx * lx + gy * ly) + + # Normalize to [-1, 1] range + max_val = max(np.abs(bevel_light).max(), 1e-8) + bevel_light = bevel_light / max_val + + # Only apply near edges + edge_proximity = np.clip(1.0 - dist / (bevel_width * 1.5), 0, 1) * hard_mask + edge_proximity += np.clip(1.0 - dist_outside / (bevel_width * 0.8), 0, 1) * (1 - hard_mask) + + bevel_light *= edge_proximity + + return bevel_light + + +def create_inner_relief(img_arr, mask, light_dir=(0.3, -0.4), strength=0.08): + """ + Detect color boundaries within the patch (where different stitch sections + meet) and add subtle height relief at those boundaries. In real embroidery, + the NASA text sits slightly above the blue fill, the chevron overlaps, etc. + """ + h, w = img_arr.shape[:2] + + # Detect edges within the patch using color gradient magnitude + # Use all 3 channels for better boundary detection + edges = np.zeros((h, w), dtype=np.float64) + for c in range(3): + gx = ndimage.sobel(img_arr[:, :, c], axis=1) + gy = ndimage.sobel(img_arr[:, :, c], axis=0) + edges += np.sqrt(gx ** 2 + gy ** 2) + edges /= 3.0 + + # Threshold to find significant color boundaries (not just stitch texture) + # Smooth to get section-level boundaries, not individual thread edges + edges_smooth = gaussian_filter(edges, sigma=1.5) + + # Normalize + edge_max = np.percentile(edges_smooth[mask > 0.5], 95) if np.any(mask > 0.5) else 1.0 + edges_norm = np.clip(edges_smooth / (edge_max + 1e-8), 0, 1) + + # Create height map: sections have flat heights, boundaries have transitions + # Use edge magnitude as a proxy for height discontinuity + height = gaussian_filter(edges_norm, sigma=2.0) * mask + + # Compute directional lighting on this height map + gx = ndimage.sobel(height, axis=1) + gy = ndimage.sobel(height, axis=0) + lx, ly = light_dir + relief = -(gx * lx + gy * ly) * strength * mask + + return relief + + +# ───────────────────────────────────────────── +# 10-11. Photographic effects +# ───────────────────────────────────────────── + +def add_vignette(img_arr, strength=0.22, radius=0.65): + """Photographic vignette - darkens corners.""" + h, w = img_arr.shape[:2] + y = np.linspace(-1, 1, h)[:, None] + x = np.linspace(-1, 1, w)[None, :] + dist = np.sqrt(x * x + y * y) + vignette = 1.0 - strength * np.clip((dist - radius) / (1.4 - radius), 0, 1) ** 1.5 + return img_arr * vignette[:, :, np.newaxis] + + +def add_film_grain(img_arr, strength=3.0): + """Photographic film grain with realistic grain size.""" + h, w = img_arr.shape[:2] + # Luminance-dependent grain (stronger in shadows) + lum = np.mean(img_arr[:, :, :3], axis=2) + grain_strength = strength * (1.0 + 0.3 * (1.0 - lum / 255.0)) + + grain = np.random.normal(0, 1, (h, w)) * grain_strength + grain = gaussian_filter(grain, sigma=0.4) + + result = img_arr + grain[:, :, np.newaxis] + return np.clip(result, 0, 255) + + +def add_depth_of_field(img_arr, mask, max_blur=1.3): + """Subtle DOF: patch center sharp, frame edges soft.""" + h, w = img_arr.shape[:2] + + y_coords, x_coords = np.where(mask > 0.5) + if len(y_coords) == 0: + return img_arr + cx, cy = int(np.mean(x_coords)), int(np.mean(y_coords)) + + y = np.arange(h)[:, None] + x = np.arange(w)[None, :] + dist = np.sqrt(((x - cx) / w * 2) ** 2 + ((y - cy) / h * 2) ** 2) + blur_t = np.clip((dist - 0.35) / 0.65, 0, 1) ** 1.5 + + blurred = np.stack([ + gaussian_filter(img_arr[:, :, c], sigma=max_blur) + for c in range(img_arr.shape[2]) + ], axis=2) + + blend = blur_t[:, :, np.newaxis] + return img_arr * (1 - blend) + blurred * blend + + +def color_grade(img_arr, warmth=0.02, contrast=1.05, saturation=1.10): + """Subtle photographic color grading.""" + result = img_arr.copy() + + # Warmth + result[:, :, 0] *= (1.0 + warmth) + result[:, :, 1] *= (1.0 + warmth * 0.2) + result[:, :, 2] *= (1.0 - warmth * 0.4) + + # S-curve contrast (gentle) + mid = 128.0 + result = mid + (result - mid) * contrast + + # Saturation boost (embroidery is vivid) + gray = np.mean(result[:, :, :3], axis=2, keepdims=True) + result[:, :, :3] = gray + (result[:, :, :3] - gray) * saturation + + return np.clip(result, 0, 255) + + +# ───────────────────────────────────────────── +# Main pipeline +# ───────────────────────────────────────────── + +def postprocess_photorealistic(input_path, output_path, + fabric_color=(35, 35, 40), + padding=55, + light_dir=(0.25, -0.35, 0.90)): + """ + Full photorealistic post-processing pipeline. + """ + print(f"Loading {input_path}...") + stitch_img = Image.open(input_path) + stitch_arr = np.array(stitch_img)[:, :, :3].astype(np.float64) + h, w = stitch_arr.shape[:2] + + print(" [1/13] Extracting patch mask...") + mask = extract_patch_mask(stitch_arr) + + print(" [2/13] Computing normal map...") + normals = compute_normal_map(stitch_arr, mask, strength=2.0) + + print(" [3/13] Estimating thread directions...") + thread_angle, anisotropy = estimate_thread_direction(stitch_arr, mask, block_size=12) + + print(" [4/13] Applying Blinn-Phong + anisotropic shading...") + lit = apply_lighting(stitch_arr, normals, mask, thread_angle, anisotropy, + light_dir=light_dir, + ambient=0.55, diffuse_strength=0.35, + specular_strength=0.12, shininess=18.0, + aniso_strength=0.08, aniso_shininess=30.0) + + print(" [5/13] Computing ambient occlusion...") + ao = compute_ambient_occlusion(stitch_arr, mask, radius=2.0, strength=0.12) + # Apply AO: darken stitch valleys + for c in range(3): + lit[:, :, c] *= (1.0 - ao) + lit = np.clip(lit, 0, 255) + + print(" [6/13] Adding per-thread micro-highlights...") + lit = add_thread_microhighlights(lit, normals, mask, thread_angle, intensity=0.05) + + # --- Canvas setup --- + canvas_w = w + padding * 2 + canvas_h = h + padding * 2 + + print(f" [7/13] Generating fabric texture ({canvas_w}x{canvas_h})...") + fabric = generate_fabric_texture(canvas_w, canvas_h, color=fabric_color) + canvas = fabric.astype(np.float64) + + # --- Pad mask to canvas size --- + mask_padded = np.zeros((canvas_h, canvas_w), dtype=np.float64) + mask_padded[padding:padding + h, padding:padding + w] = mask + + # --- Shadow --- + print(" [8/13] Creating drop shadow...") + shadow = create_patch_shadow(mask_padded, offset=(5, 6), blur_radius=10, opacity=0.50) + for c in range(3): + canvas[:, :, c] *= (1.0 - shadow * 0.8) + + # --- Composite lit patch onto canvas --- + print(" [9/13] Compositing patch onto fabric...") + for c in range(3): + patch_channel = lit[:, :, c] + canvas_region = canvas[padding:padding + h, padding:padding + w, c] + canvas[padding:padding + h, padding:padding + w, c] = ( + canvas_region * (1 - mask) + patch_channel * mask + ) + + # --- Edge bevel (thickness illusion) --- + print(" [10/13] Adding edge bevel for 3D thickness...") + bevel = create_edge_bevel(mask_padded, bevel_width=5, + light_dir=(light_dir[0], light_dir[1])) + # Apply bevel as brightness modulation + bevel_intensity = 45.0 # How strong the bevel highlight/shadow is + for c in range(3): + canvas[:, :, c] += bevel * bevel_intensity + canvas = np.clip(canvas, 0, 255) + + # --- Inner relief --- + print(" [11/13] Adding inner relief at section boundaries...") + inner_relief = create_inner_relief(stitch_arr, mask, + light_dir=(light_dir[0], light_dir[1]), + strength=0.07) + relief_padded = np.zeros((canvas_h, canvas_w), dtype=np.float64) + relief_padded[padding:padding + h, padding:padding + w] = inner_relief + for c in range(3): + canvas[:, :, c] *= (1.0 + relief_padded) + canvas = np.clip(canvas, 0, 255) + + print(" [12/13] Adding merrow edge border...") + merrow, merrow_band = create_merrow_edge(mask_padded, thickness=3) + + # Detect dominant edge color from the patch border pixels + # Sample colors from the edge region of the lit patch + edge_sample_mask = (mask > 0.3) & (mask < 0.95) + if np.any(edge_sample_mask): + edge_colors = lit[edge_sample_mask] + avg_edge = np.mean(edge_colors, axis=0) + # Make merrow edge slightly lighter than border + edge_color = np.clip(avg_edge * 1.3 + 30, 0, 255) + else: + edge_color = np.array([160, 160, 165], dtype=np.float64) + + # Apply merrow edge with bump lighting for 3D thread appearance + # Create mini-normal from merrow pattern for lit edge + merrow_gx = ndimage.sobel(merrow, axis=1) + merrow_gy = ndimage.sobel(merrow, axis=0) + merrow_light = -(merrow_gx * light_dir[0] + merrow_gy * light_dir[1]) + merrow_light = merrow_light / (np.abs(merrow_light).max() + 1e-8) * 0.3 + + for c in range(3): + edge_val = edge_color[c] * (1.0 + merrow_light) * merrow + canvas[:, :, c] = canvas[:, :, c] * (1 - merrow_band * 0.65) + edge_val * 0.65 + canvas = np.clip(canvas, 0, 255) + + # --- Photographic finishing --- + print(" [13/13] Photographic finishing (DOF, color, vignette, grain)...") + canvas = add_depth_of_field(canvas, mask_padded, max_blur=1.2) + canvas = color_grade(canvas, warmth=0.02, contrast=1.06, saturation=1.12) + canvas = add_vignette(canvas, strength=0.20, radius=0.60) + canvas = add_film_grain(canvas, strength=3.0) + + # --- Save --- + result = np.clip(canvas, 0, 255).astype(np.uint8) + output_img = Image.fromarray(result, "RGB") + output_img.save(str(output_path), quality=95) + print(f" Saved to {output_path}") + + return output_path + + +def save_debug_stages(input_path, output_dir): + """Save intermediate stages for inspection.""" + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + stitch_img = Image.open(input_path) + stitch_arr = np.array(stitch_img)[:, :, :3].astype(np.float64) + h, w = stitch_arr.shape[:2] + + mask = extract_patch_mask(stitch_arr) + normals = compute_normal_map(stitch_arr, mask, strength=2.0) + thread_angle, anisotropy = estimate_thread_direction(stitch_arr, mask, block_size=12) + + # 1. Normal map visualization (standard purple/green/blue encoding) + normal_vis = ((normals + 1.0) * 0.5 * 255).astype(np.uint8) + Image.fromarray(normal_vis, "RGB").save(str(output_dir / "01_normal_map.png")) + + # 2. Mask + mask_vis = (mask * 255).astype(np.uint8) + Image.fromarray(mask_vis, "L").save(str(output_dir / "02_mask.png")) + + # 3. Thread direction + anisotropy + thread_vis = np.zeros((h, w, 3), dtype=np.uint8) + thread_vis[:, :, 0] = ((np.cos(thread_angle) + 1) * 0.5 * 255 * mask).astype(np.uint8) + thread_vis[:, :, 1] = ((np.sin(thread_angle) + 1) * 0.5 * 255 * mask).astype(np.uint8) + thread_vis[:, :, 2] = (anisotropy * 255 * mask).astype(np.uint8) + Image.fromarray(thread_vis, "RGB").save(str(output_dir / "03_thread_direction.png")) + + # 4. Lit patch (after shading, before compositing) + lit = apply_lighting(stitch_arr, normals, mask, thread_angle, anisotropy) + lit = add_thread_microhighlights(lit, normals, mask, thread_angle) + Image.fromarray(np.clip(lit, 0, 255).astype(np.uint8), "RGB").save( + str(output_dir / "04_lit_patch.png")) + + # 5. Edge bevel visualization + mask_padded = np.zeros((h + 100, w + 100), dtype=np.float64) + mask_padded[50:50 + h, 50:50 + w] = mask + bevel = create_edge_bevel(mask_padded, bevel_width=5) + bevel_vis = ((bevel + 1) * 0.5 * 255).astype(np.uint8) + Image.fromarray(bevel_vis, "L").save(str(output_dir / "05_edge_bevel.png")) + + # 6. Shadow + shadow = create_patch_shadow(mask_padded) + shadow_vis = (shadow * 255).astype(np.uint8) + Image.fromarray(shadow_vis, "L").save(str(output_dir / "06_shadow.png")) + + print(f"Debug stages saved to {output_dir}/") + + +def create_comparison(input_path, output_path, comparison_path): + """Create side-by-side comparison image.""" + original = Image.open(input_path) + result = Image.open(output_path) + + # Make them the same height + orig_w, orig_h = original.size + res_w, res_h = result.size + + # Scale original to match result height + scale = res_h / orig_h + orig_scaled = original.resize((int(orig_w * scale), res_h), Image.LANCZOS) + + # Create comparison canvas + gap = 20 + comp_w = orig_scaled.width + res_w + gap + comp = Image.new("RGB", (comp_w, res_h + 40), (30, 30, 30)) + + # Paste images + comp.paste(orig_scaled, (0, 0)) + comp.paste(result, (orig_scaled.width + gap, 0)) + + # Add labels + draw = ImageDraw.Draw(comp) + draw.text((orig_scaled.width // 2 - 30, res_h + 5), "BEFORE", fill=(180, 180, 180)) + draw.text((orig_scaled.width + gap + res_w // 2 - 20, res_h + 5), "AFTER", fill=(180, 180, 180)) + + comp.save(str(comparison_path), quality=95) + print(f"Comparison saved to {comparison_path}") + + +def main(): + if len(sys.argv) < 2: + print(f"Usage: {sys.argv[0]} [output.png] [--debug] [--compare]") + sys.exit(1) + + input_path = Path(sys.argv[1]) + + # Find output path + positional_args = [a for a in sys.argv[2:] if not a.startswith("--")] + output_path = Path(positional_args[0]) if positional_args else \ + input_path.with_name(input_path.stem + "_photorealistic.png") + + debug = "--debug" in sys.argv + compare = "--compare" in sys.argv + + if debug: + save_debug_stages(input_path, output_path.parent / "debug") + + postprocess_photorealistic(input_path, output_path) + + if compare: + comp_path = output_path.with_name(output_path.stem + "_comparison.png") + create_comparison(input_path, output_path, comp_path) + + print("Done!") + + +if __name__ == "__main__": + main() diff --git a/worker/pipeline/svg2patch.py b/worker/pipeline/svg2patch.py new file mode 100644 index 0000000..4464c00 --- /dev/null +++ b/worker/pipeline/svg2patch.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +""" +svg2patch: Add inkstitch embroidery parameters to any SVG and render as a patch. + +Takes a clean SVG (from vtracer, manual design, etc.) and: + 1. Adds inkstitch namespace + fill parameters to all paths + 2. Sets document dimensions to mm + 3. Renders with inkstitch for realistic stitch simulation + 4. Optionally runs photorealistic post-processing +""" + +import sys +import os +import subprocess +import argparse +from pathlib import Path +from lxml import etree + +INKSTITCH_NS = "http://inkstitch.org/namespace" +SVG_NS = "http://www.w3.org/2000/svg" +INKSCAPE_NS = "http://www.inkscape.org/namespaces/inkscape" +SODIPODI_NS = "http://sodipodi.sourceforge.net/DTD/sodipodi-0.0.dtd" + +INKSTITCH_BIN = os.environ.get( + "INKSTITCH_BIN", + os.path.expanduser( + "~/Library/Application Support/org.inkscape.Inkscape" + "/config/inkscape/extensions/inkstitch.app/Contents/MacOS/inkstitch" + ), +) + +PATCH_WIDTH_MM = 80.0 + + +def add_inkstitch_params(svg_path, output_svg_path, border_color=None): + """Add inkstitch embroidery parameters to all paths in an SVG.""" + etree.register_namespace("inkstitch", INKSTITCH_NS) + etree.register_namespace("inkscape", INKSCAPE_NS) + etree.register_namespace("sodipodi", SODIPODI_NS) + + tree = etree.parse(str(svg_path)) + root = tree.getroot() + + nsmap = dict(root.nsmap) + nsmap["inkstitch"] = INKSTITCH_NS + nsmap["inkscape"] = INKSCAPE_NS + nsmap["sodipodi"] = SODIPODI_NS + new_root = etree.Element(root.tag, nsmap=nsmap) + new_root.attrib.update(root.attrib) + new_root.text = root.text + new_root.tail = root.tail + for child in root: + new_root.append(child) + root = new_root + + # get original dimensions from viewBox or width/height + viewbox = root.get("viewBox") + if viewbox: + parts = viewbox.split() + vb_w = float(parts[2]) - float(parts[0]) + vb_h = float(parts[3]) - float(parts[1]) + else: + vb_w = float(root.get("width", "100").replace("px", "").replace("mm", "")) + vb_h = float(root.get("height", "100").replace("px", "").replace("mm", "")) + + # set dimensions to mm + scale = PATCH_WIDTH_MM / vb_w + width_mm = PATCH_WIDTH_MM + height_mm = vb_h * scale + root.set("width", f"{width_mm}mm") + root.set("height", f"{height_mm}mm") + + # add namedview + existing_nv = root.find("{%s}namedview" % SODIPODI_NS) + if existing_nv is None: + nv = etree.SubElement(root, "{%s}namedview" % SODIPODI_NS) + nv.set("{%s}document-units" % INKSCAPE_NS, "mm") + + # add version metadata + existing_meta = root.find("{%s}metadata" % SVG_NS) + if existing_meta is None: + existing_meta = root.find("metadata") + if existing_meta is None: + existing_meta = etree.SubElement(root, "metadata") + version_el = existing_meta.find("{%s}inkstitch_svg_version" % INKSTITCH_NS) + if version_el is None: + version_el = etree.SubElement(existing_meta, "{%s}inkstitch_svg_version" % INKSTITCH_NS) + version_el.text = "3" + + # find all paths/shapes and add inkstitch params + all_elements = root.iter() + shape_tags = { + "{%s}path" % SVG_NS, "{%s}circle" % SVG_NS, "{%s}ellipse" % SVG_NS, + "{%s}rect" % SVG_NS, "{%s}polygon" % SVG_NS, + "path", "circle", "ellipse", "rect", "polygon", + } + + element_count = 0 + for el in all_elements: + tag = el.tag + if tag not in shape_tags: + continue + + # get fill color from style or fill attribute + style = el.get("style", "") + fill = el.get("fill", "") + + if "fill:none" in style or fill == "none": + continue + if "display:none" in style: + continue + + # convert fill attribute to style if needed + if fill and "fill:" not in style: + if style: + el.set("style", f"fill:{fill};stroke:none;{style}") + else: + el.set("style", f"fill:{fill};stroke:none") + if el.get("fill"): + del el.attrib["fill"] + elif not fill and "fill:" not in style: + continue + + # ensure stroke:none is in style + current_style = el.get("style", "") + if "stroke:" not in current_style: + el.set("style", current_style.rstrip(";") + ";stroke:none") + + # add inkstitch fill parameters + angle = (30 + element_count * 23) % 180 + el.set("{%s}fill_method" % INKSTITCH_NS, "auto_fill") + el.set("{%s}fill_underlay" % INKSTITCH_NS, "true") + el.set("{%s}fill_underlay_angle" % INKSTITCH_NS, str((angle + 90) % 360)) + el.set("{%s}angle" % INKSTITCH_NS, str(angle)) + el.set("{%s}row_spacing_mm" % INKSTITCH_NS, "0.25") + el.set("{%s}max_stitch_length_mm" % INKSTITCH_NS, "3.0") + el.set("{%s}staggers" % INKSTITCH_NS, "4") + + element_count += 1 + + # add border if requested + if border_color: + if viewbox: + parts = [float(p) for p in viewbox.split()] + origin_x, origin_y = parts[0], parts[1] + else: + origin_x, origin_y = 0.0, 0.0 + + pad = vb_w * 0.06 + border = etree.Element("rect") + border.set("x", str(origin_x - pad)) + border.set("y", str(origin_y - pad)) + border.set("width", str(vb_w + pad * 2)) + border.set("height", str(vb_h + pad * 2)) + border.set("rx", str(pad * 0.8)) + border.set("ry", str(pad * 0.8)) + border.set("style", f"fill:{border_color};stroke:none") + border.set("{%s}fill_method" % INKSTITCH_NS, "auto_fill") + border.set("{%s}fill_underlay" % INKSTITCH_NS, "true") + border.set("{%s}angle" % INKSTITCH_NS, "90") + border.set("{%s}row_spacing_mm" % INKSTITCH_NS, "0.2") + border.set("{%s}max_stitch_length_mm" % INKSTITCH_NS, "2.5") + border.set("{%s}staggers" % INKSTITCH_NS, "4") + root.insert(0, border) + + root.set("viewBox", f"{origin_x-pad} {origin_y-pad} {vb_w+pad*2} {vb_h+pad*2}") + new_scale = PATCH_WIDTH_MM / (vb_w + pad * 2) + root.set("width", f"{(vb_w + pad*2) * new_scale}mm") + root.set("height", f"{(vb_h + pad*2) * new_scale}mm") + element_count += 1 + + etree.ElementTree(root).write(str(output_svg_path), xml_declaration=True, encoding="utf-8", pretty_print=True) + print(f" {element_count} elements parameterized") + return output_svg_path + + +def render_inkstitch(svg_path, output_png): + """Render with inkstitch realistic PNG.""" + with open(output_png, "wb") as f: + result = subprocess.run( + [INKSTITCH_BIN, "--extension=png_realistic", str(svg_path)], + stdout=f, stderr=subprocess.PIPE, timeout=300, + ) + if result.returncode != 0: + print(f" inkstitch error: {result.stderr.decode()[:300]}", file=sys.stderr) + if output_png.exists(): + output_png.unlink() + return False + return output_png.exists() and output_png.stat().st_size > 0 + + +def main(): + parser = argparse.ArgumentParser(description="Add inkstitch params to SVG and render as patch") + parser.add_argument("input", help="Input SVG path") + parser.add_argument("output", nargs="?", help="Output PNG path") + parser.add_argument("-b", "--border-color", default="#0a0a14", help="Border color (default: #0a0a14)") + parser.add_argument("--no-border", action="store_true", help="Skip border") + parser.add_argument("--no-postprocess", action="store_true", help="Skip photorealistic post-processing") + args = parser.parse_args() + + input_path = Path(args.input) + output_path = Path(args.output) if args.output else input_path.with_name(input_path.stem + "_patch.png") + + border = None if args.no_border else args.border_color + + print(f"Adding inkstitch params to {input_path}...") + embroidery_svg = output_path.with_suffix(".svg") + add_inkstitch_params(input_path, embroidery_svg, border_color=border) + + print(f"Rendering with inkstitch...") + if render_inkstitch(embroidery_svg, output_path): + print(f" Stitch render: {output_path} ({output_path.stat().st_size // 1024}KB)") + else: + print(" inkstitch render failed") + return + + if not args.no_postprocess: + postprocess_script = Path(__file__).parent / "photorealistic.py" + if postprocess_script.exists(): + final_path = output_path.with_name(output_path.stem.replace("_patch", "") + "_final.png") + print(f"Post-processing...") + subprocess.run( + [sys.executable, str(postprocess_script), str(output_path), str(final_path)], + timeout=120, + ) + if final_path.exists(): + print(f" Final: {final_path}") + else: + print(" (photorealistic.py not found, skipping post-processing)") + + print("Done!") + + +if __name__ == "__main__": + main() diff --git a/worker/requirements.txt b/worker/requirements.txt new file mode 100644 index 0000000..076b92f --- /dev/null +++ b/worker/requirements.txt @@ -0,0 +1,6 @@ +redis==5.0.0 +lxml==5.2.0 +numpy==1.26.0 +scipy==1.13.0 +Pillow==10.4.0 +vtracer==0.6.15 diff --git a/worker/worker.py b/worker/worker.py new file mode 100644 index 0000000..6dd665b --- /dev/null +++ b/worker/worker.py @@ -0,0 +1,76 @@ +import os +import json +import tempfile +from pathlib import Path + +import redis + +from pipeline.convert import convert, convert_svg + +REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379") +RESULT_TTL = int(os.environ.get("RESULT_TTL", "3600")) +QUEUE_NAME = "patches" + + +def main(): + conn = redis.from_url(REDIS_URL) + conn.ping() + print(f"Worker ready, listening on queue '{QUEUE_NAME}'") + + while True: + _, payload_bytes = conn.brpop(QUEUE_NAME) + job = json.loads(payload_bytes) + job_id = job["job_id"] + print(f"Processing {job_id}") + + try: + run_pipeline( + conn, job_id, + border_color=job.get("border_color", "#0a0a14"), + color_precision=job.get("color_precision", 8), + postprocess=job.get("postprocess", True), + ) + print(f" done") + except Exception as e: + print(f" failed: {e}") + + +def run_pipeline(conn, job_id, border_color="#0a0a14", + color_precision=8, postprocess=True): + try: + input_bytes = conn.get(f"job:{job_id}:input") + ext = (conn.get(f"job:{job_id}:ext") or b"png").decode() + if not input_bytes: + raise RuntimeError("input not found in redis") + + is_svg = ext.lower() == "svg" + + with tempfile.TemporaryDirectory(prefix="p2p_") as tmpdir: + tmpdir = Path(tmpdir) + input_path = tmpdir / f"input.{ext}" + input_path.write_bytes(input_bytes) + output_path = tmpdir / "patch.png" + + if is_svg: + convert_svg(str(input_path), str(output_path), + border_color=border_color, postprocess=postprocess) + else: + convert(str(input_path), str(output_path), + border_color=border_color, + color_precision=color_precision, + postprocess=postprocess) + + result_bytes = output_path.read_bytes() + + conn.setex(f"job:{job_id}:result", RESULT_TTL, result_bytes) + conn.setex(f"job:{job_id}:status", RESULT_TTL, "complete") + conn.delete(f"job:{job_id}:input") + + except Exception as e: + conn.setex(f"job:{job_id}:status", RESULT_TTL, "failed") + conn.setex(f"job:{job_id}:error", RESULT_TTL, str(e)) + raise + + +if __name__ == "__main__": + main()