Spaces:
Sleeping
Sleeping
File size: 3,250 Bytes
2201868 52fd9c2 2201868 52fd9c2 2255b93 52fd9c2 66345ab 5479ea5 52fd9c2 fc29cbf 52fd9c2 95250f9 5cadf06 52fd9c2 95250f9 5cadf06 52fd9c2 95250f9 5cadf06 95250f9 b4d05af 95250f9 5649d80 95250f9 5649d80 2255b93 5649d80 52fd9c2 2255b93 5cadf06 5649d80 9dfc63c fb8a03b 52fd9c2 fb8a03b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import gradio as gr
import json
import torch
from torch import nn
from torchvision import models, transforms
from huggingface_hub import hf_hub_download
from PIL import Image
import requests
import os
import base64
from io import BytesIO
# Define the number of classes
num_classes = 2
# Download model from Hugging Face
def download_model():
try:
model_path = hf_hub_download(repo_id="jays009/Restnet50", filename="pytorch_model.bin")
return model_path
except Exception as e:
print(f"Error downloading model: {e}")
return None
# Load the model from Hugging Face
def load_model(model_path):
try:
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
except Exception as e:
print(f"Error loading model: {e}")
return None
# Download the model and load it
model_path = download_model()
model = load_model(model_path) if model_path else None
# Define the transformation for the input image
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
def predict(image):
try:
print(f"Received image input: {image}")
# Check if the input is a PIL Image type (Gradio automatically provides a PIL image)
if not isinstance(image, Image.Image):
return json.dumps({"error": "Invalid image format received. Please provide a valid image."})
# Apply transformations to the image
image = transform(image).unsqueeze(0)
print(f"Transformed image tensor: {image.shape}")
# Move the image to the correct device
image = image.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
# Make predictions
with torch.no_grad():
outputs = model(image)
predicted_class = torch.argmax(outputs, dim=1).item()
print(f"Prediction output: {outputs}, Predicted class: {predicted_class}")
# Return the result based on the predicted class
if predicted_class == 0:
return json.dumps({"result": "The photo you've sent is of fall army worm with problem ID 126."})
elif predicted_class == 1:
return json.dumps({"result": "The photo you've sent is of a healthy maize image."})
else:
return json.dumps({"error": "Unexpected class prediction."})
except Exception as e:
print(f"Error processing image: {e}")
return json.dumps({"error": f"Error processing image: {e}"})
# Create the Gradio interface
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil", label="Upload an image or provide a URL or local path"),
outputs=gr.Textbox(label="Prediction Result"),
live=True,
title="Maize Anomaly Detection",
description="Upload an image of maize to detect anomalies like disease or pest infestation. You can provide local paths, URLs, or base64-encoded images."
)
# Launch the Gradio interface
iface.launch(share=True, show_error=True)
|