added conv block

app.py

@@ -1,117 +1,3 @@
-# # detect.py
-# import torch
-# import torchvision.transforms as transforms
-# from torchvision.models import resnet50
-# from PIL import Image
-# import torch.nn as nn
-
-# # Define the class names - make sure these match your training classes
-# CLASS_NAMES = [
-#     "Apple___Apple_scab",
-#     "Apple___Black_rot",
-#     # Add all your class names here...
-# ]
-
-# def load_model(model_path):
-#     # Initialize the model architecture
-#     model = resnet50(pretrained=False)
-#     num_classes = len(CLASS_NAMES)
-#     model.fc = nn.Linear(model.fc.in_features, num_classes)
-    
-#     # Load the state dict
-#     state_dict = torch.load(model_path, map_location=torch.device('cpu'))
-#     model.load_state_dict(state_dict)
-#     model.eval()
-#     return model
-
-# def predict_image(image_path, model):
-#     """Predict the class of a given image"""
-#     # Define the same transform as used during training
-#     transform = transforms.Compose([
-#         transforms.Resize((224, 224)),
-#         transforms.ToTensor(),
-#     ])
-    
-#     # Load and preprocess the image
-#     image = Image.open(image_path).convert('RGB')
-#     image_tensor = transform(image).unsqueeze(0)
-    
-#     # Make prediction
-#     with torch.no_grad():
-#         outputs = model(image_tensor)
-#         _, predicted = torch.max(outputs, 1)
-        
-#     return CLASS_NAMES[predicted.item()]
-
-# # streamlit_app.py
-# import streamlit as st
-# import torch
-# import torchvision.transforms as transforms
-# from PIL import Image
-# import os
-# from detect import load_model, predict_image, CLASS_NAMES
-
-# # Set page config
-# st.set_page_config(page_title="Plant Disease Predictor", page_icon="🍃", layout="wide")
-
-# # Load the model
-# @st.cache_resource
-# def load_model_cached():
-#     model_path = 'models/leaf_disease_res50_model_epoch_10.pth'
-#     model = load_model(model_path)
-#     return model
-
-# # Load model at startup
-# model = load_model_cached()
-
-# # Streamlit app
-# st.title("Plant Disease Predictor")
-# st.write("Upload an image of a plant leaf to predict if it has a disease.")
-
-# uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
-# if uploaded_file is not None:
-#     image = Image.open(uploaded_file).convert('RGB')
-#     st.image(image, caption='Uploaded Image', use_column_width=True)
-    
-#     if st.button('Predict'):
-#         # Show prediction in progress
-#         with st.spinner('Analyzing image...'):
-#             # Save the uploaded file temporarily
-#             with open("temp_image.jpg", "wb") as f:
-#                 f.write(uploaded_file.getbuffer())
-            
-#             # Make prediction
-#             prediction = predict_image("temp_image.jpg", model)
-            
-#             # Remove temporary file
-#             os.remove("temp_image.jpg")
-            
-#             # Display result
-#             st.success(f"Prediction: {prediction}")
-            
-#             # Display confidence scores
-#             transform = transforms.Compose([
-#                 transforms.Resize((224, 224)),  # Match the training size
-#                 transforms.ToTensor(),
-#             ])
-            
-#             with torch.no_grad():
-#                 img_tensor = transform(image).unsqueeze(0)
-#                 outputs = model(img_tensor)
-#                 probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
-            
-#             # Display top 5 predictions
-#             top5_prob, top5_catid = torch.topk(probabilities, 5)
-#             st.write("Top 5 Predictions:")
-#             for i in range(top5_prob.size(0)):
-#                 st.write(f"{CLASS_NAMES[top5_catid[i]]}: {top5_prob[i].item()*100:.2f}%")
-
-# # Display list of detectable diseases
-# st.write("## List of Detectable Plant Diseases")
-# st.write("This model can detect the following plant diseases:")
-# for disease in CLASS_NAMES:
-#     st.write(f"- {disease.replace('___', ' - ')}")
 
 import gradio as gr
 import torch
@@ -121,6 +7,33 @@ import json
 import os
 # from leaf_disease_predict import ResNet9, load_model, predict_image, CLASS_NAMES
 
+class ImageClassificationBase(torch.nn.Module):
+    def validation_step(self, batch):
+        images, labels = batch
+        out = self(images)
+        loss = torch.nn.functional.cross_entropy(out, labels)
+        acc = accuracy(out, labels)
+        return {"val_loss": loss.detach(), "val_accuracy": acc}
+    
+    def validation_epoch_end(self, outputs):
+        batch_losses = [x["val_loss"] for x in outputs]
+        batch_accuracy = [x["val_accuracy"] for x in outputs]
+        epoch_loss = torch.stack(batch_losses).mean()
+        epoch_accuracy = torch.stack(batch_accuracy).mean()
+        return {"val_loss": epoch_loss, "val_accuracy": epoch_accuracy}
+    
+    def epoch_end(self, epoch, result):
+        print("Epoch [{}], last_lr: {:.5f}, train_loss: {:.4f}, val_loss: {:.4f}, val_acc: {:.4f}".format(
+            epoch, result['lrs'][-1], result['train_loss'], result['val_loss'], result['val_accuracy']))
+
+def ConvBlock(in_channels, out_channels, pool=False):
+    layers = [torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+             torch.nn.BatchNorm2d(out_channels),
+             torch.nn.ReLU(inplace=True)]
+    if pool:
+        layers.append(torch.nn.MaxPool2d(4))
+    return torch.nn.Sequential(*layers)
+
 class ResNet9(ImageClassificationBase):
     def __init__(self, in_channels, num_diseases):
         super().__init__()