Update app.py

app.py

@@ -9,12 +9,18 @@ from torchvision.models import resnet18
 import os
 import requests
 
-# UI title and instructions
+# App 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.markdown(
+    """
+    Upload a full-face image of a stroke patient.
+    The app will detect the **affected facial side** using a stroke classification model,
+    and then use the **unaffected side** to predict **pain intensity**.
+    """
+)
 st.write("🔧 App started. Preparing to download models...")
 
-# Function to download and load models with caching
+# Function to download and load models
 @st.cache_resource
 def download_models():
     model_urls = {
@@ -22,69 +28,70 @@ 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}...")
             r = requests.get(url)
             with open(filename, "wb") as f:
                 f.write(r.content)
-            st.write(f"✅ {filename} downloaded.")
+            st.success(f"✅ {filename} downloaded.")
         else:
             st.write(f"✔️ {filename} already exists.")
 
     st.write("📦 Loading models...")
 
-    # Load stroke side classification model (Keras)
+    # Load stroke side detection model (Keras)
     stroke_model = load_model("cnn_stroke_model.keras")
 
-    # Load pain intensity prediction model (PyTorch ResNet18)
+    # Load pain intensity model (PyTorch ResNet18)
     pain_model = resnet18(weights=None)
-    pain_model.fc = nn.Linear(pain_model.fc.in_features, 1)  # One output: PSPI pain score
+    pain_model.fc = nn.Linear(pain_model.fc.in_features, 1)
     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
+# Download and load models
 stroke_model, pain_model = download_models()
 
-# Define preprocessing pipeline for PyTorch pain model
+# Image preprocessing for pain model (ResNet)
 transform = transforms.Compose([
-    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
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])
 ])
 
-# File uploader widget
+# Upload image
 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")  # Ensure RGB format
+    full_image = Image.open(uploaded_file).convert("RGB")
     st.image(full_image, caption="Uploaded Full-Face Image", use_column_width=True)
 
-    # Crop left and right halves of the face
+    # Crop left and right sides
     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))
 
-    # Resize full image for stroke model input
+    # Resize 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)  # Shape: (1, H, W, C)
+    stroke_array = np.expand_dims(stroke_array, axis=0)
 
-    # Predict affected side using stroke model
+    # Predict affected side
     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
+    affected = int(np.round(stroke_pred[0][0]))  # Assuming 0 = left affected, 1 = right affected
 
-    # Select unaffected side and label sides (from patient's perspective)
-    if affected == 1:
+    # ✅ FIX: Interpret correctly based on stroke model
+    # If 0: left side affected → use right for pain
+    # If 1: right side affected → use left for pain
+    if affected == 0:
         affected_side = "left"
         unaffected_side = "right"
         unaffected_face = right_face
@@ -93,19 +100,36 @@ if uploaded_file is not None:
         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)
+    # Preprocess for pain model
+    unaffected_tensor = transform(unaffected_face).unsqueeze(0)
 
-    # Predict pain score
+    # Predict pain score (PSPI)
     st.write("📈 Predicting pain score...")
     with torch.no_grad():
         output = pain_model(unaffected_tensor)
-        pspi_score = output.item()  # Extract float value
+        pspi_score = output.item()
 
-    # Display results
-    st.subheader("🧾 Prediction Results")
+    # Display prediction 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"**🧭 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]}`")
+    st.write(f"**📊 Stroke model raw output:** `{stroke_pred[0][0]:.4f}`")
+
+    # Pain scale explanation
+    st.markdown(
+        """
+        ---
+        ### ℹ️ **What is PSPI?**
+        The **Prkachin and Solomon Pain Intensity (PSPI)** score ranges from **0 to 6**:
+
+        - `0`: No pain
+        - `1–2`: Mild pain
+        - `3–4`: Moderate pain
+        - `5–6`: Severe pain
+
+        This score is computed from facial action units such as brow lowering, eye closure, and cheek raising.
+        """
+    )