Spaces:
Sleeping
Sleeping
File size: 1,101 Bytes
aa1c1e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
import uuid
import os
from models.utils import save_mask_as_png
def run_prompt_mode(image, prompt_text, sam_wrapper, dino_wrapper, save_dir="outputs"):
"""
Process image using DINO + SAM based on user prompt.
:param image: Input image (BGR format)
:param prompt_text: Text prompt from user
:param sam_wrapper: Initialized SamWrapper
:param dino_wrapper: Initialized DinoWrapper
:param clip_wrapper: Optional CLIP for verification
:param save_dir: Directory to save results
:return: List of saved file paths
"""
saved_paths = []
os.makedirs(save_dir, exist_ok=True)
# 1. Detect boxes with DINO
boxes = dino_wrapper.detect(image, prompt_text)
# 2. Segment with SAM for each box
for i, box in enumerate(boxes):
mask = sam_wrapper.predict_with_box(image, box)
if mask is None:
continue # safety check
filename = f"{save_dir}/{uuid.uuid4().hex[:8]}_{i}_{prompt_text.replace(' ', '_')}.png"
save_mask_as_png(image, mask, filename)
saved_paths.append(filename)
return saved_paths
|