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()