Add project files

pipeline/sam3_client.py

@@ -0,0 +1,103 @@
+"""
+Bridge from the main app (torch17_new, Python 3.8) to SAM3 (sam3_worker.py,
+run under sam2_env's Python 3.10 -- see config.py's SAM3 section for why
+this has to be a subprocess rather than an in-process import).
+
+One subprocess invocation can batch multiple (image, prompt) jobs -- e.g.
+"flower" on the upload and "vase" on both the upload and the matched
+template -- so the ~6s model load only happens once per request, not once
+per job.
+"""
+
+import json
+import logging
+import os
+import shutil
+import subprocess
+import uuid
+
+import cv2
+import numpy as np
+
+import config
+
+logger = logging.getLogger(__name__)
+
+
+class Sam3Error(RuntimeError):
+    pass
+
+
+def run_jobs(images: dict, jobs: list, workdir: str) -> dict:
+    """images: {image_key: bgr_ndarray}. jobs: [{"image": key, "prompt": str,
+    "threshold": float}, ...]. Returns {(image_key, prompt): [{"mask": bool
+    ndarray, "score": float, "box": [x1,y1,x2,y2] or None}, ...]}.
+
+    Raises Sam3Error on any failure (missing token, subprocess crash,
+    timeout, malformed response) -- callers treat this the same as any
+    other soft/optional-feature failure (caught, logged, reported back as
+    an error string, never taking the whole request down)."""
+    if not config.HF_TOKEN:
+        raise Sam3Error(
+            "No HF_TOKEN found (checked environment and .env) -- SAM3's "
+            "weights are gated on Hugging Face and can't be downloaded "
+            "without an access-granted token."
+        )
+
+    request_dir = os.path.join(workdir, f"sam3_{uuid.uuid4().hex[:8]}")
+    os.makedirs(request_dir, exist_ok=True)
+    try:
+        image_paths = {}
+        for key, bgr in images.items():
+            path = os.path.join(request_dir, f"{key}.png")
+            cv2.imwrite(path, bgr)
+            image_paths[key] = path
+
+        output_dir = os.path.join(request_dir, "out")
+        request = {"images": image_paths, "jobs": jobs, "output_dir": output_dir}
+        request_path = os.path.join(request_dir, "request.json")
+        with open(request_path, "w") as f:
+            json.dump(request, f)
+
+        env = dict(os.environ)
+        env["HF_TOKEN"] = config.HF_TOKEN
+
+        proc = subprocess.run(
+            [config.SAM3_PYTHON_BIN, config.SAM3_WORKER_SCRIPT, "--request", request_path],
+            capture_output=True, text=True, timeout=config.SAM3_TIMEOUT_SECONDS, env=env,
+        )
+
+        response_path = os.path.join(output_dir, "response.json")
+        if not os.path.isfile(response_path):
+            raise Sam3Error(
+                f"SAM3 worker produced no response (exit {proc.returncode}): "
+                f"{proc.stderr[-2000:] if proc.stderr else '(no stderr)'}"
+            )
+
+        with open(response_path) as f:
+            response = json.load(f)
+
+        if response.get("error"):
+            raise Sam3Error(f"SAM3 worker failed: {response['error'][:2000]}")
+
+        out = {}
+        for entry in response["results"]:
+            key = (entry["image"], entry["prompt"])
+            instances = []
+            for inst in entry["instances"]:
+                mask_path = os.path.join(output_dir, inst["mask_file"])
+                mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
+                if mask is None:
+                    continue
+                instances.append({
+                    "mask": mask > 127,
+                    "score": inst["score"],
+                    "box": inst.get("box"),
+                })
+            out[key] = instances
+        return out
+
+    except subprocess.TimeoutExpired as e:
+        raise Sam3Error(f"SAM3 worker timed out after {config.SAM3_TIMEOUT_SECONDS}s") from e
+    finally:
+        shutil.rmtree(request_dir, ignore_errors=True)