File size: 3,189 Bytes
97cf772 | 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 torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import gradio as gr
# -------------------------
# Classes (first version) - exact order from training
# -------------------------
classes = [
'Pepper__bell___Bacterial_spot',
'Pepper__bell___healthy',
'Potato___Early_blight',
'Potato___Late_blight',
'Potato___healthy',
'Tomato_Bacterial_spot',
'Tomato_Early_blight',
'Tomato_Late_blight',
'Tomato_Leaf_Mold',
'Tomato_Septoria_leaf_spot',
'Tomato_Spider_mites_Two_spotted_spider_mite',
'Tomato__Target_Spot',
'Tomato__Tomato_YellowLeaf__Curl_Virus',
'Tomato__Tomato_mosaic_virus',
'Tomato_healthy'
]
# Optional Hebrew translations (adjust as needed)
class_map_hebrew = {
'Pepper__bell___Bacterial_spot': "ืคืืคื ืืชืืง - ืืชื ืืืืืงื",
'Pepper__bell___healthy': "ืคืืคื ืืชืืง - ืืจืื",
'Potato___Early_blight': "ืชืคืื ืืืื - ืจืืงืืื ืืืงืื",
'Potato___Late_blight': "ืชืคืื ืืืื - ืจืืงืืื ืืืืืจ",
'Potato___healthy': "ืชืคืื ืืืื - ืืจืื",
'Tomato_Bacterial_spot': "ืขืืื ืืื - ืืชื ืืืืืงื",
'Tomato_Early_blight': "ืขืืื ืืื - ืจืืงืืื ืืืงืื",
'Tomato_Late_blight': "ืขืืื ืืื - ืจืืงืืื ืืืืืจ",
'Tomato_Leaf_Mold': "ืขืืื ืืื - ืขืืืฉ ืขืื",
'Tomato_Septoria_leaf_spot': "ืขืืื ืืื - ืืชื ืขืื Septoria",
'Tomato_Spider_mites_Two_spotted_spider_mite': "ืขืืื ืืื - ืงืจืืืช ืฉื ื ืืชืืื",
'Tomato__Target_Spot': "ืขืืื ืืื - ืืชื ืืืจื",
'Tomato__Tomato_YellowLeaf__Curl_Virus': "ืขืืื ืืื - ืฆืืืช ืขืื / ืืืจืืก ืืชืืืืฆืืช",
'Tomato__Tomato_mosaic_virus': "ืขืืื ืืื - ืืืจืืก ืคืกืืคืก",
'Tomato_healthy': "ืขืืื ืืื - ืืจืืื"
}
# -------------------------
# Build ResNet18 and load checkpoint
# -------------------------
model = models.resnet18(pretrained=False)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, len(classes))
# Load the full checkpoint (including fc)
checkpoint = torch.load("resnet18_15class.pth", map_location="cpu")
model.load_state_dict(checkpoint)
model.eval()
# -------------------------
# Image preprocessing
# -------------------------
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(img):
image = transform(img).unsqueeze(0)
with torch.no_grad():
outputs = model(image)
_, predicted = torch.max(outputs, 1)
english = classes[predicted.item()]
hebrew = class_map_hebrew[english]
return f"{english} ({hebrew})"
# -------------------------
# Gradio interface
# -------------------------
demo = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs="text",
title="GreenGuard - Leaf Disease Detection"
)
if __name__ == "__main__":
demo.launch()
|