paintplus: vendor the app source and rename from EditmaskwithAI

Bring the full EditmaskwithAI application into the repo under paintplus/ (429 files) so the service is self-contained — the installer copies the vendored source to ~/docker/paintplus/src instead of cloning at runtime. Rename to PaintPlus (service + branding; app logic untouched): - services/editmaskwithai.sh -> services/paintplus.sh (register_service paintplus, install_paintplus, ~/docker/paintplus, Caddy paintplus:8000, Authelia option preserved) - container names -> paintplus across docker-compose*.yml; dev network -> paintplus-network - browser <title> -> "PaintPlus - AI Image Editor"; README heading -> PaintPlus with upstream provenance note - README utilities table: editmaskwithai -> paintplus Backend/frontend code (help strings referencing the old container name, the ai_photo_edit.db filename) is intentionally left as-is to avoid touching application logic. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Nb2vJ8W7bHKx1JXVvpCraH
2026-06-26 05:48:43 +00:00
parent b4e8ba2a79
commit 084922afaa
431 changed files with 87396 additions and 77 deletions
@@ -0,0 +1,184 @@
+"""
+SAM (Segment Anything Model) service.
+
+Auto-downloads the ViT-B checkpoint (~375 MB) on first use.
+Caches the loaded model in memory; re-uses predictor across calls.
+
+Prediction API:
+  predict_points(image_bytes, points, labels) -> mask_bytes (PNG, white=selected)
+  points: list of (x, y) in original image pixels
+  labels: list of 1 (include) or 0 (exclude), same length as points
+"""
+
+import asyncio
+import io
+import os
+import urllib.request
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+from PIL import Image
+
+# ── Model download ────────────────────────────────────────────────────────────
+
+SAM_DIR      = Path("/app/data/models/sam")
+SAM_FILENAME = "sam_vit_b_01ec64.pth"
+SAM_URL      = f"https://dl.fbaipublicfiles.com/segment_anything/{SAM_FILENAME}"
+SAM_PATH     = SAM_DIR / SAM_FILENAME
+
+
+class SamInstallState(str, Enum):
+    idle        = "idle"
+    downloading = "downloading"
+    done        = "done"
+    failed      = "failed"
+
+
+@dataclass
+class SamInstallStatus:
+    state:    SamInstallState = SamInstallState.idle
+    progress: int             = 0
+    message:  str             = ""
+    error:    str             = ""
+
+
+_install_status = SamInstallStatus()
+_install_lock   = asyncio.Lock()
+
+
+def get_install_status() -> dict:
+    s = _install_status
+    return {"state": s.state.value, "progress": s.progress,
+            "message": s.message, "error": s.error}
+
+
+def sam_model_available() -> bool:
+    return SAM_PATH.exists() and SAM_PATH.stat().st_size > 100_000_000
+
+
+async def ensure_sam_installed() -> bool:
+    """Download SAM ViT-B checkpoint if not present. Returns True on success."""
+    global _install_status
+
+    if sam_model_available():
+        _install_status = SamInstallStatus(state=SamInstallState.done, progress=100,
+                                           message="SAM model ready.")
+        return True
+
+    async with _install_lock:
+        if sam_model_available():
+            _install_status = SamInstallStatus(state=SamInstallState.done, progress=100,
+                                               message="SAM model ready.")
+            return True
+
+        if _install_status.state == SamInstallState.downloading:
+            return False
+
+        try:
+            SAM_DIR.mkdir(parents=True, exist_ok=True)
+            _install_status = SamInstallStatus(
+                state=SamInstallState.downloading, progress=0,
+                message="Downloading SAM ViT-B model (~375 MB)…",
+            )
+
+            def _download():
+                def _progress(count, block, total):
+                    if total > 0:
+                        _install_status.progress = min(99, int(count * block * 99 / total))
+                tmp = SAM_PATH.with_suffix(".tmp")
+                urllib.request.urlretrieve(SAM_URL, tmp, _progress)
+                tmp.rename(SAM_PATH)
+
+            loop = asyncio.get_event_loop()
+            await loop.run_in_executor(None, _download)
+
+            _install_status = SamInstallStatus(state=SamInstallState.done, progress=100,
+                                               message="SAM model ready.")
+            return True
+
+        except Exception as exc:
+            _install_status = SamInstallStatus(
+                state=SamInstallState.failed, error=str(exc),
+                message="SAM download failed.",
+            )
+            print(f"[sam] Download failed: {exc}")
+            return False
+
+
+# ── Model cache ───────────────────────────────────────────────────────────────
+
+_predictor = None
+_predictor_lock = asyncio.Lock()
+
+
+def _load_predictor():
+    """Load SAM model and return a SamPredictor. Called in thread pool."""
+    global _predictor
+    if _predictor is not None:
+        return _predictor
+
+    import torch
+    from segment_anything import sam_model_registry, SamPredictor
+
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+
+    print(f"[sam] Loading SAM ViT-B on {device}…")
+    sam = sam_model_registry["vit_b"](checkpoint=str(SAM_PATH))
+    sam.to(device)
+    _predictor = SamPredictor(sam)
+    print("[sam] Model loaded.")
+    return _predictor
+
+
+# ── Prediction ────────────────────────────────────────────────────────────────
+
+def _predict_sync(image_bytes: bytes,
+                  points: list[tuple[int, int]],
+                  labels: list[int]) -> bytes:
+    """
+    Run SAM prediction synchronously (call via run_in_executor).
+    Returns PNG bytes: white = selected, black = background.
+    """
+    predictor = _load_predictor()
+
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    img_array = np.array(image)
+
+    predictor.set_image(img_array)
+
+    pt_array  = np.array(points, dtype=np.float32)   # [[x, y], ...]
+    lbl_array = np.array(labels, dtype=np.int32)      # [1=fg, 0=bg, ...]
+
+    masks, scores, _ = predictor.predict(
+        point_coords=pt_array,
+        point_labels=lbl_array,
+        multimask_output=True,
+    )
+
+    # Pick the highest-confidence mask
+    best = masks[int(np.argmax(scores))]  # bool array H×W
+
+    mask_img = Image.fromarray((best * 255).astype(np.uint8), mode="L")
+    buf = io.BytesIO()
+    mask_img.save(buf, format="PNG")
+    return buf.getvalue()
+
+
+async def predict_points(image_bytes: bytes,
+                         points: list[tuple[int, int]],
+                         labels: list[int]) -> bytes:
+    """Async wrapper for SAM point prediction."""
+    if not sam_model_available():
+        ok = await ensure_sam_installed()
+        if not ok:
+            raise RuntimeError("SAM model not available.")
+    loop = asyncio.get_event_loop()
+    return await loop.run_in_executor(None, _predict_sync, image_bytes, points, labels)