Spaces:
Sleeping
Sleeping
File size: 1,281 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 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import cv2
import os
from models.sam import SamWrapper
from models.dino import DinoWrapper
from pipeline.prompt_mode import run_prompt_mode
def main_prompt():
"""
Run prompt-based segmentation using DINO + SAM.
"""
# --- Ensure necessary folders exist
os.makedirs("assets", exist_ok=True)
os.makedirs("outputs", exist_ok=True)
image_path = "assets/fir_tree.jpg"
output_dir = "outputs"
# prompt = "a vegetable with purple leaves" # You can change this prompt
# prompt = "a long light green vegetable"
# prompt = "a round fruit with a red rough surface"
# prompt = "a long green cucumber"
# prompt = "a long green cucumber"
# prompt = "a yellow car indoor"
prompt = "a high fir tree"
# prompt = "a car top view"
image = cv2.imread(image_path)
if image is None:
print("Failed to load image.")
return
sam = SamWrapper(
model_type="vit_b",
checkpoint_path="weights/sam_vit_b_01ec64.pth"
)
dino = DinoWrapper()
results = run_prompt_mode(image, prompt, sam, dino, save_dir=output_dir)
print(f"\nProcessed {len(results)} objects (prompt mode):")
for path in results:
print(f"→ Saved at {path}")
if __name__ == "__main__":
main_prompt()
|