Update src/siglip/classifier.py

src/siglip/classifier.py

@@ -1,6 +1,7 @@
 # Necessary imports
 import sys
 from typing import Dict
+from PIL import Image
 import torch
 from transformers import AutoModel, AutoProcessor
 import gradio as gr
@@ -16,7 +17,9 @@ model = AutoModel.from_pretrained(model_id).eval().to("cpu")
 processor = AutoProcessor.from_pretrained(model_id)
 
 
-def ZeroShotImageClassification(image_input, candidate_labels):
+def ZeroShotImageClassification(
+    image_input: Image.Image, candidate_labels: str
+) -> Dict[str, float]:
     """
     Performs zero-shot classification on the given image input and candidate labels.