Update app.py

app.py

@@ -1,23 +1,28 @@
 import torch
 import streamlit as st
+import torchvision.models as models
 import numpy as np
 import cv2
 from PIL import Image
 import torchvision.transforms as transforms
 
-# Load models
 @st.cache_resource
 def load_models():
-    # Load classification model (corrected filename)
-    classification_model = torch.load("classifier.pt", map_location=torch.device("cpu"))
+    # Load classification model architecture
+    classification_model = models.resnet152(pretrained=False)  # Use the same architecture as before
+    num_ftrs = classification_model.fc.in_features
+    classification_model.fc = torch.nn.Linear(num_ftrs, 5)  # Assuming 5 DR stages
+    classification_checkpoint = torch.load("classifier.pt", map_location=torch.device("cpu"))
+    classification_model.load_state_dict(classification_checkpoint["model_state_dict"])
     classification_model.eval()
-    
+
     # Load segmentation model
     segmentation_model = torch.load("best_unet_model.pth", map_location=torch.device("cpu"))
     segmentation_model.eval()
-    
+
     return classification_model, segmentation_model
 
+
 classification_model, segmentation_model = load_models()
 
 # Define preprocessing function for classification