Update app.py

app.py

@@ -1,8 +1,8 @@
 import streamlit as st
 import torch
 import torch.nn as nn
-from torchvision import transforms
 from torchvision.models import resnet18
+from torchvision import transforms
 from PIL import Image
 import numpy as np
 from tensorflow.keras.models import load_model
@@ -15,26 +15,26 @@ st.markdown("Upload a full-face image. The system will detect the affected side
 
 @st.cache_resource
 def download_models():
-    model_files = {
+    model_urls = {
         "cnn_stroke_model.keras": "https://huggingface.co/AdhamQQ/cnn_stroke_model/resolve/main/cnn_stroke_model.keras",
         "left_side_pain_classifier.pth": "https://huggingface.co/AdhamQQ/cnn_stroke_model/resolve/main/left_side_pain_classifier.pth",
-        "right_side_pain_classifier.pth": "https://huggingface.co/AdhamQQ/cnn_stroke_model/resolve/main/right_side_pain_classifier.pth"
+        "right_side_pain_classifier.pth": "https://huggingface.co/AdhamQQ/cnn_stroke_model/resolve/main/right_side_pain_classifier.pth",
     }
 
-    for fname, url in model_files.items():
-        if not os.path.exists(fname) or os.path.getsize(fname) < 10000:
-            st.write(f"📥 Downloading {fname}...")
+    for filename, url in model_urls.items():
+        if not os.path.exists(filename) or os.path.getsize(filename) < 10000:
+            st.write(f"📥 Downloading {filename}...")
             try:
                 r = requests.get(url)
                 r.raise_for_status()
-                with open(fname, "wb") as f:
+                with open(filename, "wb") as f:
                     f.write(r.content)
-                st.success(f"✅ {fname} downloaded.")
+                st.success(f"✅ {filename} downloaded.")
             except Exception as e:
-                st.error(f"❌ Could not download {fname}: {e}")
+                st.error(f"❌ Could not download {filename}: {e}")
                 st.stop()
         else:
-            st.write(f"✔️ {fname} already exists.")
+            st.write(f"✔️ {filename} already exists.")
 
     # Haar cascade
     haar_path = "haarcascade_frontalface_default.xml"
@@ -51,9 +51,18 @@ def download_models():
         st.error(f"❌ Failed to load stroke model: {e}")
         st.stop()
 
+    # Define correct PainModel class
+    class PainModel(nn.Module):
+        def __init__(self):
+            super(PainModel, self).__init__()
+            self.convnet = resnet18(weights=None)
+            self.convnet.fc = nn.Linear(self.convnet.fc.in_features, 1)
+
+        def forward(self, x):
+            return self.convnet(x)
+
     def load_pain_model(path):
-        model = resnet18(weights=None)
-        model.fc = nn.Linear(model.fc.in_features, 1)
+        model = PainModel()
         model.load_state_dict(torch.load(path, map_location="cpu"))
         model.eval()
         return model
@@ -72,7 +81,8 @@ stroke_model, left_model, right_model = download_models()
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])
 ])
 
 uploaded_file = st.file_uploader("📂 Upload a full-face image", type=["jpg", "jpeg", "png"])