updated

app.py

@@ -119,7 +119,62 @@ import torchvision.transforms as transforms
 from PIL import Image
 import json
 import os
-from leaf_disease_predict import ResNet9, load_model, predict_image, CLASS_NAMES
+# from leaf_disease_predict import ResNet9, load_model, predict_image, CLASS_NAMES
+
+class ResNet9(ImageClassificationBase):
+    def __init__(self, in_channels, num_diseases):
+        super().__init__()
+        
+        self.conv1 = ConvBlock(in_channels, 64)
+        self.conv2 = ConvBlock(64, 128, pool=True)
+        self.res1 = torch.nn.Sequential(ConvBlock(128, 128), ConvBlock(128, 128))
+        
+        self.conv3 = ConvBlock(128, 256, pool=True)
+        self.conv4 = ConvBlock(256, 512, pool=True)
+        self.res2 = torch.nn.Sequential(ConvBlock(512, 512), ConvBlock(512, 512))
+        
+        self.classifier = torch.nn.Sequential(torch.nn.MaxPool2d(4),
+                                       torch.nn.Flatten(),
+                                       torch.nn.Linear(512, num_diseases))
+        
+    def forward(self, xb):
+        out = self.conv1(xb)
+        out = self.conv2(out)
+        out = self.res1(out) + out
+        out = self.conv3(out)
+        out = self.conv4(out)
+        out = self.res2(out) + out
+        out = self.classifier(out)
+        return out
+
+CLASS_NAMES = [
+    'Apple___Apple_scab', 'Apple___Black_rot', 'Apple___Cedar_apple_rust', 'Apple___healthy',
+    'Blueberry___healthy', 'Cherry_(including_sour)___Powdery_mildew', 'Cherry_(including_sour)___healthy',
+    
+]
+
+def predict_image(image_path, model):
+    transform = transforms.Compose([
+        transforms.Resize((256, 256)),
+        transforms.ToTensor(),
+    ])
+    
+    img = Image.open(image_path).convert('RGB')
+    img_tensor = transform(img).unsqueeze(0)
+    
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        _, predicted = torch.max(outputs, 1)
+        
+    return CLASS_NAMES[predicted.item()]
+
+
+def load_model(model_path):
+    model = torch.load(model_path, map_location=torch.device('cpu'))
+    model.eval()
+    return model
+
+
 
 # Load the model
 model_path = 'models/leaf_disease_res50_model_epoch_10.pth'