141 lines
4.7 KiB
Python
141 lines
4.7 KiB
Python
"""
|
|
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()
|