Update app.py

app.py

@@ -1,49 +1,50 @@
 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):
+# ----------------------
+# Inference function
+# ----------------------
+def classify_weather(image_file):
     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.")
+        # Open the uploaded file
+        image = Image.open(image_file).convert("RGB")
 
-        # preprocess as batch
+        # Preprocess
         inputs = processor(images=[image], return_tensors="pt")
 
-        # inference
+        # 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()
+            logits = outputs.logits.squeeze()
+            probs = torch.softmax(logits, dim=-1).tolist()
             labels = [model.config.id2label[i] for i in range(len(probs))]
-            output_dict = dict(zip(labels, probs))
 
-        return output_dict
+            # Return label -> probability dictionary
+            return dict(zip(labels, probs))
 
     except Exception as e:
-        return {"Error": str(e)}
-
+        # Safe fallback if something unexpected happens
+        return {"Error": 1.0}
 
+# ----------------------
 # Gradio interface
+# ----------------------
 iface = gr.Interface(
     fn=classify_weather,
-    inputs=gr.Image(type="numpy"),
+    inputs=gr.File(file_types=[".jpg", ".png"]),  # Accept uploaded files
     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 Space with error reporting
 if __name__ == "__main__":
     iface.launch(show_error=True)