Update app.py

app.py

@@ -1,49 +1,49 @@
-
-import gradio as gr
-import numpy as np
-from PIL import Image
-import torch
-from transformers import AutoImageProcessor, AutoModelForImageClassification
-
-# ----------------------
-# Load model + processor
-# ----------------------
-processor = AutoImageProcessor.from_pretrained("prithivMLmods/Weather-Image-Classification")
-model = AutoModelForImageClassification.from_pretrained("prithivMLmods/Weather-Image-Classification")
-
-# ----------------------
-# Inference function
-# ----------------------
-def classify_weather(image_input):
-    # Only NumPy array supported for Gradio input
-    if isinstance(image_input, np.ndarray):
-        image = Image.fromarray(image_input.astype('uint8')).convert("RGB")
-    else:
-        raise TypeError("Only NumPy array input is supported for this Gradio interface.")
-
-    # Preprocess
-    inputs = processor(images=image, return_tensors="pt")
-
-    # Inference
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-        predicted_class_id = logits.argmax(-1).item()
-        predicted_label = model.config.id2label[predicted_class_id]
-
-    return predicted_label
-
-# ----------------------
-# Gradio interface
-# ----------------------
-iface = gr.Interface(
-    fn=classify_weather,
-    inputs=gr.Image(type="numpy"),  # NumPy array input
-    outputs=gr.Label(num_top_classes=5, label="Weather Condition"),
-    title="Weather Image Classification",
-    description="Upload an image to classify the weather condition (sun, rain, snow, fog, or clouds)."
-)
-
-# Launch the app
-if __name__ == "__main__":
-    iface.launch()
+import gradio as gr
+import numpy as np
+from PIL import Image
+import torch
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+# Load model + processor
+processor = AutoImageProcessor.from_pretrained("prithivMLmods/Weather-Image-Classification")
+model = AutoModelForImageClassification.from_pretrained("prithivMLmods/Weather-Image-Classification")
+
+def classify_weather(image_input):
+    try:
+        if isinstance(image_input, np.ndarray):
+            image = Image.fromarray(image_input.astype("uint8")).convert("RGB")
+        else:
+            raise TypeError("Only NumPy array input is supported.")
+
+        # preprocess as batch
+        inputs = processor(images=[image], return_tensors="pt")
+
+        # inference
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits
+            predicted_class_id = logits.argmax(-1).item()
+            predicted_label = model.config.id2label[predicted_class_id]
+
+            # optional: return probabilities for Label(num_top_classes=5)
+            probs = torch.softmax(logits, dim=-1).squeeze().tolist()
+            labels = [model.config.id2label[i] for i in range(len(probs))]
+            output_dict = dict(zip(labels, probs))
+
+        return output_dict
+
+    except Exception as e:
+        return {"Error": str(e)}
+
+
+# Gradio interface
+iface = gr.Interface(
+    fn=classify_weather,
+    inputs=gr.Image(type="numpy"),
+    outputs=gr.Label(num_top_classes=5, label="Weather Condition"),
+    title="Weather Image Classification",
+    description="Upload an image to classify the weather condition (sun, rain, snow, fog, or clouds)."
+)
+
+if __name__ == "__main__":
+    iface.launch(show_error=True)