Update README.md

README.md

@@ -290,8 +290,7 @@ The model categorizes images into the following 126 classes:
 
 ```python
 import gradio as gr
-from transformers import AutoImageProcessor
-from transformers import SiglipForImageClassification
+from transformers import AutoImageProcessor, SiglipForImageClassification
 from transformers.image_utils import load_image
 from PIL import Image
 import torch
@@ -302,14 +301,21 @@ model = SiglipForImageClassification.from_pretrained(model_name)
 processor = AutoImageProcessor.from_pretrained(model_name)
 
 def sketch_classification(image):
-    \"\"\"Predicts the sketch category for an input image.\"\"\n    image = Image.fromarray(image).convert(\"RGB\")
-    inputs = processor(images=image, return_tensors=\"pt\")
+    """Predicts the sketch category for an input image."""
+    # Convert the input numpy array to a PIL Image and ensure it has 3 channels (RGB)
+    image = Image.fromarray(image).convert("RGB")
     
+    # Process the image and prepare it for the model
+    inputs = processor(images=image, return_tensors="pt")
+    
+    # Perform inference without gradient calculation
     with torch.no_grad():
         outputs = model(**inputs)
         logits = outputs.logits
+        # Convert logits to probabilities using softmax
         probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
     
+    # Mapping from indices to corresponding sketch category labels
     labels = {
         "0": "aircraft_carrier", "1": "alarm_clock", "2": "ant", "3": "anvil", "4": "asparagus",
         "5": "axe", "6": "banana", "7": "basket", "8": "bathtub", "9": "bear",
@@ -339,20 +345,21 @@ def sketch_classification(image):
         "122": "vase", "123": "watermelon", "124": "whale", "125": "zebra"
     }
     
+    # Create a dictionary mapping each label to its predicted probability (rounded)
     predictions = {labels[str(i)]: round(probs[i], 3) for i in range(len(probs))}
     return predictions
 
 # Create Gradio interface
 iface = gr.Interface(
     fn=sketch_classification,
-    inputs=gr.Image(type=\"numpy\"),
-    outputs=gr.Label(label=\"Prediction Scores\"),
-    title=\"Sketch-126-DomainNet Classification\",
-    description=\"Upload a sketch to classify it into one of 126 categories.\"
+    inputs=gr.Image(type="numpy"),
+    outputs=gr.Label(label="Prediction Scores"),
+    title="Sketch-126-DomainNet Classification",
+    description="Upload a sketch to classify it into one of 126 categories."
 )
 
 # Launch the app
-if __name__ == \"__main__\":
+if __name__ == "__main__":
     iface.launch()
 ```