Add project files

sam3_worker.py

@@ -0,0 +1,140 @@
+"""
+Standalone SAM3 concept-segmentation worker.
+
+Runs under a SEPARATE Python environment (sam2_env, Python 3.10) from the
+main Vase Matcher app (torch17_new, Python 3.8) -- transformers>=5.5.0
+(required for Sam3Model/Sam3Processor) itself requires Python>=3.10, so
+this can't be imported in-process by the Flask app. Instead it's invoked
+as a one-shot subprocess per request (see pipeline/sam3_client.py) with a
+JSON request file describing which prompts to run against which images,
+and writes a JSON response + one PNG mask per detected instance.
+
+Loaded in 4-bit (NF4) via bitsandbytes -- empirically ~700MB resident /
+~1.9GB peak during inference on an 8GB card, vs. several GB unquantized,
+and a single model load handles every job in the request (one process
+per /api/count_flowers call, not per prompt/image).
+
+Usage:
+    python sam3_worker.py --request <request.json>
+
+Request JSON:
+    {
+      "images": {"<image_key>": "<path to bgr PNG/JPG>", ...},
+      "jobs": [{"image": "<image_key>", "prompt": "flower", "threshold": 0.5}, ...],
+      "output_dir": "<dir to write masks + response.json into>"
+    }
+
+Response JSON (written to <output_dir>/response.json):
+    {
+      "results": [
+        {"image": "...", "prompt": "...", "instances": [
+            {"mask_file": "...", "score": 0.73, "box": [x1,y1,x2,y2]}, ...
+        ]}, ...
+      ],
+      "error": null   # or a string on failure
+    }
+"""
+
+import argparse
+import json
+import os
+import sys
+import traceback
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+HF_MODEL_NAME = "facebook/sam3"
+
+
+def load_model():
+    import torch
+    from transformers import BitsAndBytesConfig, Sam3Model, Sam3Processor
+
+    token = os.environ.get("HF_TOKEN")
+    processor = Sam3Processor.from_pretrained(HF_MODEL_NAME, token=token)
+    quant_config = BitsAndBytesConfig(
+        load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4"
+    )
+    model = Sam3Model.from_pretrained(
+        HF_MODEL_NAME, token=token, quantization_config=quant_config,
+        device_map="cuda:0", low_cpu_mem_usage=True,
+    )
+    model.eval()
+    return model, processor
+
+
+def run_job(model, processor, pil_img, prompt, threshold):
+    import torch
+
+    inputs = processor(images=pil_img, text=prompt, return_tensors="pt").to("cuda:0")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    results = processor.post_process_instance_segmentation(
+        outputs, threshold=threshold, mask_threshold=0.5,
+        target_sizes=[pil_img.size[::-1]],
+    )[0]
+
+    masks = results["masks"].cpu().numpy()
+    scores = results["scores"].cpu().numpy().tolist()
+    boxes = results["boxes"].cpu().numpy().tolist() if "boxes" in results else [None] * len(masks)
+    return masks, scores, boxes
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--request", required=True)
+    args = parser.parse_args()
+
+    with open(args.request) as f:
+        request = json.load(f)
+
+    output_dir = request["output_dir"]
+    os.makedirs(output_dir, exist_ok=True)
+    response = {"results": [], "error": None}
+
+    try:
+        import cv2
+        import numpy as np
+        from PIL import Image
+
+        model, processor = load_model()
+
+        loaded_images = {}
+        for key, path in request["images"].items():
+            bgr = cv2.imread(path, cv2.IMREAD_COLOR)
+            if bgr is None:
+                raise ValueError(f"could not read image {path!r} for key {key!r}")
+            loaded_images[key] = Image.fromarray(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
+
+        for job_idx, job in enumerate(request["jobs"]):
+            image_key = job["image"]
+            prompt = job["prompt"]
+            threshold = job.get("threshold", 0.5)
+            pil_img = loaded_images[image_key]
+
+            masks, scores, boxes = run_job(model, processor, pil_img, prompt, threshold)
+
+            instances = []
+            for i, (m, score, box) in enumerate(zip(masks, scores, boxes)):
+                m_bin = (m > 0.5).astype("uint8") * 255
+                mask_file = f"{image_key}_{prompt.replace(' ', '_')}_{job_idx}_{i}.png"
+                cv2.imwrite(os.path.join(output_dir, mask_file), m_bin)
+                instances.append({"mask_file": mask_file, "score": round(float(score), 4), "box": box})
+
+            response["results"].append({
+                "image": image_key, "prompt": prompt, "instances": instances,
+            })
+
+    except Exception as e:
+        response["error"] = f"{e}\n{traceback.format_exc()}"
+
+    with open(os.path.join(output_dir, "response.json"), "w") as f:
+        json.dump(response, f)
+
+    if response["error"]:
+        print(response["error"], file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()