Upload folder using huggingface_hub

README.md

@@ -7,7 +7,7 @@ sdk: gradio
 sdk_version: 6.6.0
 app_file: app.py
 pinned: false
-short_description: Florence-2 tagging for SAM 3 auto-discovery
+short_description: Florence-2 multi-task tagging for SAM 3 auto-discovery
 ---
 
-`api_autotag(image, max_tags)` -> JSON {labels: [...], tags: [{name, count}]} via Florence-2 <OD>.
+`api_autotag(image, max_tags)` -> JSON `{labels: [...], tags: [{name, count}], tasks: [...]}` via Florence-2 `<OD>`, `<DENSE_REGION_CAPTION>`, and `<MORE_DETAILED_CAPTION>`.

app.py

@@ -1,6 +1,14 @@
-"""Florence-2 object tagger (ZeroGPU): proposes class names from an image."""
+"""Florence-2 scene tagger (ZeroGPU): proposes class names from an image.
+
+The first implementation used only the Florence ``<OD>`` task, which is too
+sparse for street-level panoptic work because object detection misses broad
+surface/stuff classes. This endpoint now combines object detection, dense
+region captions, and detailed caption text into one open-vocabulary concept
+list for SAM3.
+"""
 import os
 from collections import Counter
+import re
 
 import gradio as gr
 import spaces
@@ -14,27 +22,102 @@ processor = AutoProcessor.from_pretrained(MODEL_ID)
 model = Florence2ForConditionalGeneration.from_pretrained(MODEL_ID)
 model.eval()
 
+TASKS = ("<OD>", "<DENSE_REGION_CAPTION>", "<MORE_DETAILED_CAPTION>")
+DROP_WORDS = {
+    "a", "an", "the", "this", "that", "these", "those", "there", "here",
+    "image", "photo", "picture", "view", "scene", "background", "foreground",
+    "left", "right", "top", "bottom", "front", "back", "side", "area", "part",
+    "visible", "large", "small", "several", "multiple", "many", "some",
+}
+CAPTION_SPLIT = re.compile(
+    r"[,.;:]|\\bwith\\b|\\band\\b|\\bnext to\\b|\\bin front of\\b|\\bbehind\\b|\\bon\\b|\\balong\\b|\\bnear\\b",
+    flags=re.IGNORECASE,
+)
+
+
+def _clean_label(value: str) -> str:
+    text = str(value).strip().lower()
+    text = re.sub(r"[_/\\-]+", " ", text)
+    text = re.sub(r"[^a-z0-9\\s]+", " ", text)
+    words = [w for w in text.split() if w and w not in DROP_WORDS]
+    if not words:
+        return ""
+    if len(words) > 5:
+        words = words[-5:]
+    return " ".join(words)
+
+
+def _labels_from_caption(text: str) -> list[str]:
+    labels: list[str] = []
+    for chunk in CAPTION_SPLIT.split(str(text)):
+        clean = _clean_label(chunk)
+        if not clean:
+            continue
+        words = clean.split()
+        # Prefer compact noun-like endings from descriptive chunks while keeping
+        # short labels intact. This is intentionally generic, not a fixed class
+        # list for one dataset.
+        for candidate in (" ".join(words[-3:]), " ".join(words[-2:]), words[-1]):
+            candidate = _clean_label(candidate)
+            if candidate and len(candidate) > 2:
+                labels.append(candidate)
+                break
+    return labels
+
+
+def _extract_labels(parsed) -> list[str]:
+    labels: list[str] = []
+
+    def walk(obj, key_hint: str = ""):
+        if isinstance(obj, dict):
+            for key, value in obj.items():
+                k = str(key).lower()
+                if k in {"label", "labels", "caption", "captions", "text", "description", "descriptions"}:
+                    walk(value, k)
+                else:
+                    walk(value, key_hint)
+        elif isinstance(obj, (list, tuple)):
+            for item in obj:
+                walk(item, key_hint)
+        elif isinstance(obj, str):
+            if key_hint in {"label", "labels"}:
+                clean = _clean_label(obj)
+                if clean:
+                    labels.append(clean)
+            else:
+                labels.extend(_labels_from_caption(obj))
+
+    walk(parsed)
+    return labels
+
+
+def _run_florence_task(image, task: str, num_beams: int) -> dict:
+    device = "cuda"
+    model.to(device)
+    inputs = processor(text=task, images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        gen = model.generate(**inputs, max_new_tokens=1536, num_beams=int(num_beams))
+    text = processor.batch_decode(gen, skip_special_tokens=False)[0]
+    parsed = processor.post_process_generation(text, task=task, image_size=image.size)
+    return {"task": task, "text": text, "parsed": parsed, "labels": _extract_labels(parsed)}
+
 
 @spaces.GPU(duration=120)
 def api_autotag(image, max_tags, num_beams=3):
     if image is None:
         return {"error": "no image provided"}
     image = image.convert("RGB")
-    device = "cuda"
-    model.to(device)
-    inputs = processor(text="<OD>", images=image, return_tensors="pt").to(device)
-    with torch.no_grad():
-        gen = model.generate(**inputs, max_new_tokens=1024, num_beams=int(num_beams))
-    text = processor.batch_decode(gen, skip_special_tokens=False)[0]
-    parsed = processor.post_process_generation(text, task="<OD>", image_size=image.size)
-    labels = parsed.get("<OD>", {}).get("labels", [])
-    counts = Counter(str(l).strip().lower() for l in labels if str(l).strip())
+    task_results = [_run_florence_task(image, task, int(num_beams)) for task in TASKS]
+    labels = []
+    for result in task_results:
+        labels.extend(result["labels"])
+    counts = Counter(label for label in labels if label)
     tags = [{"name": n, "count": c} for n, c in counts.most_common(int(max_tags))]
-    return {"model": MODEL_ID, "tags": tags, "labels": [t["name"] for t in tags]}
+    return {"model": MODEL_ID, "tasks": task_results, "tags": tags, "labels": [t["name"] for t in tags]}
 
 
 with gr.Blocks(title="SAM3 AutoTag") as demo:
-    gr.Markdown("# Florence-2 AutoTag\nUpload an image; returns detected object class names.")
+    gr.Markdown("# Florence-2 AutoTag\nUpload an image; returns multi-task scene class names for SAM3.")
     with gr.Row():
         inp = gr.Image(type="pil", label="Image")
         out = gr.JSON(label="Tags")