WeatherAI / app.py
pyedward's picture
Rename Weather.py to app.py
b1e7cc8 verified
raw
history blame
1.64 kB
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()