104 lines
3.6 KiB
Python
104 lines
3.6 KiB
Python
"""
|
|
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)
|