import gradio as gr
import torch
from torch import nn
from torchvision import models, transforms
from huggingface_hub import hf_hub_download
from PIL import Image

num_classes = 2  

def download_model():
    model_path = hf_hub_download(repo_id="jays009/Restnet50", filename="pytorch_model.bin")
    return model_path

def load_model(model_path):
    model = models.resnet50(pretrained=False)  
    model.fc = nn.Linear(model.fc.in_features, num_classes)  
    model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))  
    model.eval() 
    return model

model_path = download_model()  
model = load_model(model_path)


transform = transforms.Compose([
    transforms.Resize(256),  # Resize the image to 256x256
    transforms.CenterCrop(224),  # Crop the image to 224x224
    transforms.ToTensor(),  # Convert the image to a Tensor
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),  
])

def predict(image):
    
    image = transform(image).unsqueeze(0)  
    image = image.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))  
    
    with torch.no_grad():
        outputs = model(image)  
        predicted_class = torch.argmax(outputs, dim=1).item()  

    
    if predicted_class == 0:
        return "The photo you've sent is of fall army worm with problem ID 126."
    elif predicted_class == 1:
        return "The photo you've sent is of a healthy wheat image."
    else:
        return "Unexpected class prediction."

iface = gr.Interface(
    fn=predict,  
    inputs=gr.Image(type="pil"),  
    outputs=gr.Textbox(),  
    live=True,  
    title="Maize Anomaly Detection",
    description="Upload an image of maize to detect anomalies like disease or pest infestation."
)

iface.launch()