Update app.py

app.py

@@ -9,10 +9,12 @@ from torchvision.models import resnet18
 import os
 import requests
 
-st.title("Stroke Patient Pain Intensity Detector")
+# UI title and instructions
+st.title("🧠 Stroke Patient Pain Intensity Detector")
 st.markdown("Upload a full-face image of a stroke patient. The app will detect the affected side and predict pain intensity using the unaffected side.")
 st.write("🔧 App started. Preparing to download models...")
 
+# Function to download and load models with caching
 @st.cache_resource
 def download_models():
     model_urls = {
@@ -20,6 +22,7 @@ def download_models():
         "pain_model.pth": "https://huggingface.co/AdhamQQ/cnn_stroke_model/resolve/main/pain_model.pth"
     }
 
+    # Download models if not already downloaded
     for filename, url in model_urls.items():
         if not os.path.exists(filename):
             st.write(f"📥 Downloading {filename}...")
@@ -31,70 +34,78 @@ def download_models():
             st.write(f"✔️ {filename} already exists.")
 
     st.write("📦 Loading models...")
-    stroke_model = load_model("cnn_stroke_model.keras")
 
+    # Load stroke side classification model (Keras)
+    stroke_model = load_model("cnn_stroke_model.keras")
 
+    # Load pain intensity prediction model (PyTorch ResNet18)
     pain_model = resnet18(weights=None)
-    pain_model.fc = nn.Linear(pain_model.fc.in_features, 1)
+    pain_model.fc = nn.Linear(pain_model.fc.in_features, 1)  # One output: PSPI pain score
     pain_model.load_state_dict(torch.load("pain_model.pth", map_location=torch.device("cpu")))
     pain_model.eval()
 
-
     return stroke_model, pain_model
 
+# Load models
 stroke_model, pain_model = download_models()
 
+# Define preprocessing pipeline for PyTorch pain model
 transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                         std=[0.229, 0.224, 0.225])
+    transforms.Resize((224, 224)),  # Resize to ResNet18 input
+    transforms.ToTensor(),          # Convert to tensor
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalize using ImageNet stats
 ])
 
-uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
+# File uploader widget
+uploaded_file = st.file_uploader("📂 Choose a full-face image", type=["jpg", "jpeg", "png"])
 
+# If an image is uploaded, start processing
 if uploaded_file is not None:
     st.write("📷 Image uploaded. Processing...")
-    full_image = Image.open(uploaded_file).convert("RGB")
+    full_image = Image.open(uploaded_file).convert("RGB")  # Ensure RGB format
     st.image(full_image, caption="Uploaded Full-Face Image", use_column_width=True)
 
+    # Crop left and right halves of the face
     w, h = full_image.size
     mid = w // 2
     left_face = full_image.crop((0, 0, mid, h))
     right_face = full_image.crop((mid, 0, w, h))
 
-    # 🔍 Automatically resize image based on model input
+    # Resize full image for stroke model input
     _, H, W, C = stroke_model.input_shape
     st.write(f"🔍 Resizing uploaded image to: ({H}, {W}) for stroke model")
-
     stroke_input = full_image.resize((W, H))
     stroke_array = np.array(stroke_input).astype("float32") / 255.0
-    stroke_array = np.expand_dims(stroke_array, axis=0)
-
-st.write("🧠 Running stroke model prediction...")
-stroke_pred = stroke_model.predict(stroke_array)
-affected = int(np.round(stroke_pred[0][0]))  # 1 = left affected, 0 = right affected
-
-# Face POV logic: patient's left/right
-if affected == 1:
-    affected_side = "left"
-    unaffected_side = "right"
-    unaffected_face = right_face
-else:
-    affected_side = "right"
-    unaffected_side = "left"
-    unaffected_face = left_face
-
-unaffected_tensor = transform(unaffected_face).unsqueeze(0)
-
-st.write("📈 Predicting pain score...")
-with torch.no_grad():
-    output = pain_model(unaffected_tensor)
-    pspi_score = output.item()
-
-st.subheader("Prediction Results")
-st.image(unaffected_face, caption="Unaffected Side Used for Pain Detection", width=300)
-st.write(f"**Affected side (face POV):** {affected_side}")
-st.write(f"**Unaffected side (face POV):** {unaffected_side}")
-st.write(f"**Predicted PSPI Pain Score:** {pspi_score:.3f}")
-st.write(f"**Stroke model raw output:** {stroke_pred[0][0]}")
+    stroke_array = np.expand_dims(stroke_array, axis=0)  # Shape: (1, H, W, C)
+
+    # Predict affected side using stroke model
+    st.write("🧠 Running stroke model prediction...")
+    stroke_pred = stroke_model.predict(stroke_array)
+    affected = int(np.round(stroke_pred[0][0]))  # Output: 1 = left affected, 0 = right affected
+
+    # Select unaffected side and label sides (from patient's perspective)
+    if affected == 1:
+        affected_side = "left"
+        unaffected_side = "right"
+        unaffected_face = right_face
+    else:
+        affected_side = "right"
+        unaffected_side = "left"
+        unaffected_face = left_face
+
+    # Preprocess the unaffected side for the PyTorch model
+    unaffected_tensor = transform(unaffected_face).unsqueeze(0)  # Shape: (1, 3, 224, 224)
+
+    # Predict pain score
+    st.write("📈 Predicting pain score...")
+    with torch.no_grad():
+        output = pain_model(unaffected_tensor)
+        pspi_score = output.item()  # Extract float value
+
+    # Display results
+    st.subheader("🧾 Prediction Results")
+    st.image(unaffected_face, caption="Unaffected Side Used for Pain Detection", width=300)
+    st.write(f"**🧭 Affected side (face POV):** `{affected_side}`")
+    st.write(f"**✅ Unaffected side (face POV):** `{unaffected_side}`")
+    st.write(f"**🎯 Predicted PSPI Pain Score:** `{pspi_score:.3f}`")
+    st.write(f"**📊 Stroke model raw output:** `{stroke_pred[0][0]}`")