File size: 2,765 Bytes
cc91aff |
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 |
import torch
from torchvision import models, transforms
from PIL import Image
import gradio as gr
## Define the device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load the trained model
def load_model():
model = models.resnet50(pretrained=False)
num_classes = 4 # Update based on your rice disease classes
model.fc = torch.nn.Sequential(
torch.nn.Linear(model.fc.in_features, 256),
torch.nn.ReLU(),
torch.nn.Linear(256, num_classes)
)
model.load_state_dict(torch.load(r"/kaggle/input/rice_epoch8/pytorch/default/1/best_model_epoch_8.pth", map_location=device), strict=False)
model = model.to(device)
model.eval()
return model
# Define preprocessing steps
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# Prediction function
def predict(image):
# Ensure image is in RGB
image = image.convert("RGB")
input_tensor = transform(image).unsqueeze(0).to(device)
# Perform inference
with torch.no_grad():
outputs = model(input_tensor)
_, predicted_class = torch.max(outputs, 1)
# Map predicted class index to actual labels
class_names = ["Brown Spot", "Healthy", "Leaf Blast", "Neck Blast"]
predicted_label = class_names[predicted_class.item()]
# Calculate confidence scores
probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]
confidence = probabilities[predicted_class.item()].item()
return f"Predicted Disease: {predicted_label}\nConfidence: {confidence*100:.2f}%"
# Load the model globally
model = load_model()
# Create Gradio interface
def launch_interface():
# Create a Gradio interface
iface = gr.Interface(
theme="Subh775/orchid_candy",
fn=predict,
inputs=gr.Image(type="pil", label="Upload Rice Leaf Image"),
outputs=gr.Textbox(label="Prediction Results"),
title="Rice Disease Classification",
description="Upload a rice leaf image to detect disease type",
examples=[
["https://doa.gov.lk/wp-content/uploads/2020/06/brownspot3-1024x683.jpg"],
["https://arkansascrops.uada.edu/posts/crops/rice/images/Fig%206%20Rice%20leaf%20blast%20coalesced%20lesions.png"],
["https://th.bing.com/th/id/OIP._5ejX_5Z-M0cO5c2QUmPlwHaE7?w=280&h=187&c=7&r=0&o=5&dpr=1.1&pid=1.7"],
["https://www.weknowrice.com/wp-content/uploads/2022/11/how-to-grow-rice.jpeg"],
],
allow_flagging="never"
)
return iface
# Launch the interface
if __name__ == "__main__":
interface = launch_interface()
interface.launch(share=True) |