Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import torch
 import torch.nn as nn
 from torchvision import transforms
+from torchvision.models import resnet18
 from PIL import Image
 import numpy as np
 from tensorflow.keras.models import load_model
@@ -14,123 +15,90 @@ st.markdown("Upload a full-face image. The system will detect the affected side
 
 @st.cache_resource
 def download_models():
-    # 🧬 Ensure LFS-tracked files are downloaded
-    st.write("🔄 Checking and pulling Git LFS files...")
-    os.system("git lfs pull")
-
-    model_urls = {
+    model_files = {
         "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 filename, url in model_urls.items():
-        if not os.path.exists(filename) or os.path.getsize(filename) < 1000:
-            st.write(f"📥 Downloading {filename}...")
+    for fname, url in model_files.items():
+        if not os.path.exists(fname) or os.path.getsize(fname) < 10000:
+            st.write(f"📥 Downloading {fname}...")
             try:
-                with requests.get(url, stream=True) as r:
-                    r.raise_for_status()
-                    with open(filename, "wb") as f:
-                        for chunk in r.iter_content(chunk_size=8192):
-                            if chunk:
-                                f.write(chunk)
-                st.success(f"{filename} downloaded.")
+                r = requests.get(url)
+                r.raise_for_status()
+                with open(fname, "wb") as f:
+                    f.write(r.content)
+                st.success(f"✅ {fname} downloaded.")
             except Exception as e:
-                st.error(f"❌ Failed to download {filename}: {e}")
+                st.error(f"❌ Could not download {fname}: {e}")
                 st.stop()
         else:
-            st.write(f"✔️ {filename} already exists.")
+            st.write(f"✔️ {fname} already exists.")
 
-    # ✅ Haar cascade for face detection
+    # Haar cascade
     haar_path = "haarcascade_frontalface_default.xml"
     haar_url = "https://raw.githubusercontent.com/opencv/opencv/master/data/haarcascades/haarcascade_frontalface_default.xml"
     if not os.path.exists(haar_path):
-        st.write("📥 Downloading Haar Cascade...")
         r = requests.get(haar_url)
         with open(haar_path, "wb") as f:
             f.write(r.content)
 
-    if not os.path.exists(haar_path):
-        st.error("❌ Haar Cascade not found.")
-        st.stop()
-
-    # ✅ Check and load stroke model
-    model_path = "cnn_stroke_model.keras"
-    if not os.path.exists(model_path) or os.path.getsize(model_path) < 10000:
-        st.error("❌ Model file missing or too small. Likely not downloaded via Git LFS.")
-        st.stop()
-
+    # Load Keras model
     try:
-        stroke_model = load_model(model_path)
+        stroke_model = load_model("cnn_stroke_model.keras")
     except Exception as e:
         st.error(f"❌ Failed to load stroke model: {e}")
         st.stop()
 
-    # ✅ Define and load PyTorch pain models
-    class PainModel(nn.Module):
-        def __init__(self):
-            super(PainModel, self).__init__()
-            from torchvision.models import resnet18, ResNet18_Weights
-            self.convnet = resnet18(weights=ResNet18_Weights.DEFAULT)
-            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.load_state_dict(torch.load(path, map_location="cpu"))
+        model.eval()
+        return model
 
     try:
-        left_model = PainModel()
-        right_model = PainModel()
-        left_model.load_state_dict(torch.load("left_side_pain_classifier.pth", map_location="cpu"))
-        right_model.load_state_dict(torch.load("right_side_pain_classifier.pth", map_location="cpu"))
-        left_model.eval()
-        right_model.eval()
+        left_model = load_pain_model("left_side_pain_classifier.pth")
+        right_model = load_pain_model("right_side_pain_classifier.pth")
     except Exception as e:
         st.error(f"❌ Error loading PyTorch pain models: {e}")
         st.stop()
 
     return stroke_model, left_model, right_model
 
-# ✅ Load models
 stroke_model, left_model, right_model = download_models()
 
-# ✅ Preprocessing transform
 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])
 ])
 
-# ✅ File uploader
 uploaded_file = st.file_uploader("📂 Upload a full-face image", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
     full_image = Image.open(uploaded_file).convert("RGB")
     st.image(full_image, caption="📸 Uploaded Full-Face Image", use_column_width=True)
 
-    # ✅ Face detection
-    cv_image = np.array(full_image)
-    gray = cv2.cvtColor(cv_image, cv2.COLOR_RGB2GRAY)
+    # Face detection
+    np_img = np.array(full_image)
+    gray = cv2.cvtColor(np_img, cv2.COLOR_RGB2GRAY)
     face_cascade = cv2.CascadeClassifier("haarcascade_frontalface_default.xml")
+    faces = face_cascade.detectMultiScale(gray, 1.1, 5)
 
-    if face_cascade.empty():
-        st.error("❌ Haar cascade failed to load.")
-        st.stop()
-
-    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
     if len(faces) == 0:
-        st.error("❌ No face detected. Please upload a clear image of a face.")
+        st.error("❌ No face detected. Please upload a clearer image.")
         st.stop()
 
-    # ✅ Stroke prediction
+    # Stroke side prediction
     w, h = full_image.size
     mid = w // 2
     left_half = full_image.crop((0, 0, mid, h))
     right_half = full_image.crop((mid, 0, w, h))
 
     stroke_input = full_image.resize((224, 224))
-    stroke_array = np.array(stroke_input).astype("float32") / 255.0
-    stroke_array = np.expand_dims(stroke_array, axis=0)
+    stroke_array = np.expand_dims(np.array(stroke_input) / 255.0, axis=0)
 
     st.write("🧠 Predicting affected side...")
     stroke_pred = stroke_model.predict(stroke_array)
@@ -139,7 +107,6 @@ if uploaded_file is not None:
     unaffected_face = right_half if affected == 0 else left_half
     selected_model = right_model if affected == 0 else left_model
 
-    # ✅ Pain prediction
     st.write("📈 Predicting pain...")
     with torch.no_grad():
         tensor = transform(unaffected_face).unsqueeze(0)
@@ -147,7 +114,6 @@ if uploaded_file is not None:
         prob = torch.sigmoid(output).item()
         label = 1 if prob > 0.5 else 0
 
-    # ✅ Display results
     st.subheader("🔍 Prediction Result")
     st.image(unaffected_face, caption="🧍 Unaffected Side Used", width=300)
     st.write(f"🧭 Affected Side: **{'left' if affected == 0 else 'right'}**")