Added siglip multiple res models

app.py

@@ -11,7 +11,7 @@ BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 DEFAULT_LABELS = "gun, knife, cigarette, phone"
 
 
-def run_dfine_classify(image, dfine_threshold, dfine_model_choice, siglip_threshold, labels_text):
+def run_dfine_classify(image, dfine_threshold, dfine_model_choice, classifier_choice, siglip_threshold, labels_text):
     """D-FINE first, then classify crops with SigLIP.
     Returns (group_crop_gallery, known_crop_gallery, status_message).
     """
@@ -24,6 +24,7 @@ def run_dfine_classify(image, dfine_threshold, dfine_model_choice, siglip_thresh
 
     dfine_model = dfine_model_choice.strip().lower() if dfine_model_choice else "medium-obj2coco"
     conf_thresh = float(siglip_threshold)
+    classifier = classifier_choice.strip() if classifier_choice else "siglip-224"
 
     group_crops, known_crops, status = run_single_image(
         image,
@@ -34,7 +35,7 @@ def run_dfine_classify(image, dfine_threshold, dfine_model_choice, siglip_thresh
         min_side=24,
         crop_dedup_iou=0.4,
         min_display_conf=conf_thresh,
-        classifier="siglip",
+        classifier=classifier,
         labels=labels,
     )
 
@@ -73,6 +74,12 @@ with gr.Blocks(title="Small Object Detection") as app:
                 label="D-FINE model",
             )
 
+            classifier_dropdown = gr.Dropdown(
+                choices=["siglip-224", "siglip-256", "siglip-384"],
+                value="siglip-224",
+                label="Classifier model",
+            )
+
             dfine_threshold_slider = gr.Slider(
                 minimum=0.05,
                 maximum=0.5,
@@ -137,7 +144,7 @@ with gr.Blocks(title="Small Object Detection") as app:
 
     btn_dfine.click(
         fn=run_dfine_classify,
-        inputs=[inp_dfine, dfine_threshold_slider, dfine_model_radio, siglip_threshold_slider, labels_input],
+        inputs=[inp_dfine, dfine_threshold_slider, dfine_model_radio, classifier_dropdown, siglip_threshold_slider, labels_input],
         outputs=[out_gallery_dfine, out_gallery_known, out_status_dfine],
         concurrency_limit=1,
     )

dfine_jina_pipeline.py

@@ -516,7 +516,7 @@ DFINE_MODEL_IDS = {
     "large-obj2coco": "ustc-community/dfine-large-obj2coco-e25",
 }
 
-CLASSIFIER_CHOICES = ["jina", "siglip", "siglip2_onnx"]
+CLASSIFIER_CHOICES = ["jina", "siglip-224", "siglip-256", "siglip-384", "siglip2_onnx"]
 
 
 def _load_classifier(classifier_name, device, refs_dir=None, labels=None):
@@ -529,9 +529,11 @@ def _load_classifier(classifier_name, device, refs_dir=None, labels=None):
         ref_labels, ref_embs = build_refs(jina_encoder, refs_dir, TRUNCATE_DIM, 0.3, batch_size=16)
         return ("jina_wrapped", jina_encoder, ref_labels, ref_embs)
 
-    if classifier_name == "siglip":
-        from siglip_zeroshot import SigLIPClassifier
-        clf = SigLIPClassifier(device)
+    if classifier_name.startswith("siglip-"):
+        from siglip_zeroshot import SigLIPClassifier, SIGLIP_MODELS
+        if classifier_name not in SIGLIP_MODELS:
+            raise ValueError(f"Unknown SigLIP model: {classifier_name}. Choose from {list(SIGLIP_MODELS.keys())}")
+        clf = SigLIPClassifier(device, model_key=classifier_name)
         clf.build_refs(refs_dir=refs_dir, labels=labels)
         return clf
 
@@ -566,7 +568,7 @@ def run_single_image(
     crop_dedup_iou=0.35,
     squarify=True,
     min_display_conf=None,
-    classifier="siglip",
+    classifier="siglip-224",
     labels=None,
 ):
     """

siglip_zeroshot.py

@@ -11,17 +11,25 @@ import numpy as np
 import torch
 from transformers import SiglipModel, AutoProcessor
 
+SIGLIP_MODELS = {
+    "siglip-224": "google/siglip-base-patch16-224",
+    "siglip-256": "google/siglip-base-patch16-256",
+    "siglip-384": "google/siglip-base-patch16-384",
+}
+
 class SigLIPClassifier:
     """Zero-shot crop classifier using SigLIP (PyTorch)."""
 
-    def __init__(self, device="cuda"):
-        print("[*] Loading SigLIP (google/siglip-base-patch16-224)...")
+    def __init__(self, device="cuda", model_key="siglip-224"):
+        model_id = SIGLIP_MODELS.get(model_key, model_key)
+        print(f"[*] Loading SigLIP ({model_id})...")
         t0 = time.perf_counter()
 
         self.device = device
-        self.model = SiglipModel.from_pretrained("google/siglip-base-patch16-224")
+        self.model_key = model_key
+        self.model = SiglipModel.from_pretrained(model_id)
         self.model = self.model.to(device).eval()
-        self.processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224")
+        self.processor = AutoProcessor.from_pretrained(model_id)
 
         self.labels = []