File size: 3,712 Bytes
cc91aff 45eff52 cc91aff c797bd7 cc91aff 45eff52 cc91aff 45eff52 cc91aff 45eff52 cc91aff 45eff52 cc91aff 45eff52 cc91aff 0597a57 cc91aff 45eff52 cc91aff 45eff52 |
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 95 96 97 98 99 100 101 102 103 104 105 106 107 |
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("best_model_epoch_43.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])
])
# Color mapping for labels
label_colors = {
"Brown Spot": "#b2ff00",
"Healthy": "#2ecc71",
"Leaf Blast": "#ff00d4",
"Neck Blast": "#ffd100"
}
# Function to get color based on confidence
def get_confidence_color(confidence):
if confidence < 0.25:
return "#e74c3c" # Red
elif confidence < 0.50:
return "#f39c12"
elif confidence < 0.75:
return "#00b9ff" # Yellow
else:
return "#13ff00" # Green
# Updated 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()
# Generate styled output
label_color = label_colors.get(predicted_label, "#FFFFFF") # Default White if not found
confidence_color = get_confidence_color(confidence)
result = f"<div style='color:{label_color}; font-size:30px; font-weight:bold;'>{predicted_label}</div>"
result += f"<div style='color:{confidence_color}; font-size:25px; font-weight:bold;'>Confidence: {confidence*100:.2f}%</div>"
return result
# Updated Gradio Interface
def launch_interface():
# Create a Gradio interface
iface = gr.Interface(
theme=gr.themes.Citrus(
primary_hue="emerald",
neutral_hue="slate"
),
fn=predict,
inputs=gr.Image(type="pil", label="Upload Rice Leaf Image"),
outputs=gr.HTML(label="Prediction Results"),
title="<span style='color: #00fff7; font-size:40px; font-weight: bold;'>Rice Disease Classification</span>",
description="<span style='color: lightblue; font-size:26px;'>Upload a rice leaf image to detect its condition (Brown Spot, Healthy, Leaf Blast, or Neck Blast)</span>",
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"]
],
allow_flagging="never"
)
return iface
# Load the model globally
model = load_model()
# Launch the interface
if __name__ == "__main__":
interface = launch_interface()
interface.launch(share=True)
|