fixed the always one class prediction issue

.gitignore

@@ -0,0 +1,8 @@
+# folders in the root directory
+/models/*
+/data/*
+
+# files in the root directory
+_app.py
+_debug.py
+_howto.ipynb

app.py

@@ -1,109 +1,56 @@
-import gradio as gr
-import tensorflow as tf
-from huggingface_hub import hf_hub_download
-from datasets import load_dataset
-from PIL import Image
-import numpy as np 
-from tensorflow.keras.preprocessing.image import img_to_array
-from modelbuilder import capsnet_custom_objects
+from base import *
 
-#-------CONSTANTS-------#
-TARGET_SIZE = (256, 256)  # target size for masked images
-CLASS_LABELS =  ["COVID", "Lung_Opacity", "Normal", "Viral Pneumonia"]
 # ------------------------------------------------------------
-# 1️⃣ Load the models from Hugging Face Hub
-# capsnet for disease classification and GAN for lung segmentation/masking
+# 3️⃣ Define calback function for Gradio Interface
 # ------------------------------------------------------------
-gan_model_path = hf_hub_download(
-    repo_id="valste/lung-segmentation-gan",
-    filename="model.keras"
-)
-model_gan = tf.keras.models.load_model(gan_model_path, compile=False)
 
-capsnet_model_path = hf_hub_download(
-    repo_id="valste/capsnet-4class-lung-disease-classifier",
-    filename="model.keras"
-)
-model_capsnet = tf.keras.models.load_model(capsnet_model_path, custom_objects=capsnet_custom_objects, compile=False)
-# ------------------------------------------------------------
-# 2️⃣ Load sample X-ray images from your dataset
-# ------------------------------------------------------------
-dataset = load_dataset("valste/lung-disease-xrays", data_dir="external_xrays_299x299", split="train")
-
-demo_images = []
-for example in dataset:
-    if "image" in example:
-        img = example["image"]
-        demo_images.append(img)
-    elif "path" in example:  # If dataset stores file paths
-        img = Image.open(example["path"])
-        demo_images.append(img)
-
-# ------------------------------------------------------------
-# 3️⃣ Define preprocessing and inference function
-# ------------------------------------------------------------
-def preprocess_image(img: Image.Image):
-    # --- 1) Prepare grayscale image for GAN (1 channel) ---
-    img = img.convert("L")  # force grayscale; result shape when array: (H, W)
-    img = img.resize(TARGET_SIZE, Image.BILINEAR)  # resize to target: (width, height)
-
-    img_array = np.array(img, dtype=np.float32) / 255.0  # to float32 in [0,1], shape: (H, W)
-
-    gan_input = np.expand_dims(img_array, axis=-1)  # add channel dim for GAN -> (H, W, 1)
-    gan_input = np.expand_dims(gan_input, axis=0)   # add batch dim for GAN -> (1, H, W, 1)
-
-    # --- 2) Run segmentation GAN to get lung mask ---
-    prediction = model_gan.predict(gan_input)  # expected output e.g. (1, H, W, 1)
-
-    lung_prob = prediction[0, :, :, 0]  # remove batch + channel -> prob map, shape: (H, W)
-    mask = (lung_prob > 0.5).astype(np.float32)  # threshold to binary mask (0 or 1), shape: (H, W)
-
-    # --- 3) Apply mask on grayscale image ---
-    masked_gray = img_array * mask  # keep lung region, zero other pixels, shape: (H, W)
+def _predict_and_display(img_path: str):
+    if not img_path:
+        # return empty values to clear outputs
+        return {}, None
+    filename_out, masked_vis, scores = predict(img_path)
+    return scores, masked_vis
 
-    # --- 4) Convert masked grayscale to 3 channels for CapsNet ---
-    masked_rgb = np.stack([masked_gray] * 3, axis=-1)  # replicate into 3 channels -> (H, W, 3)
+def _clear_outputs():
+    return gr.update(value=None), gr.update(value=None)
 
-    # --- 5) Prepare two versions: one for model, one for display ---
-    x = np.expand_dims(masked_rgb, axis=0).astype(np.float32)  # model input: (1, H, W, 3)
+default_image_label = "Select or upload an image to classify" #default label for img_input
 
-    masked_vis = (masked_rgb * 255).astype(np.uint8)  # for Gradio image output (uint8 0–255)
+def _change_input_label(img_path: str | None, default_image_label=default_image_label) -> gr.update:
+    if img_path:
+        fname = os.path.basename(img_path)
+        return gr.update(label= f"Image to classify: {fname}")
+    return gr.update(label=default_image_label)
 
-    return x, masked_vis  # ready for model_capsnet + visualization
+# ------------------------------------------------------------
     
 
-def predict(image: Image.Image):
-    # x: model input, masked_vis: image to show in UI
-    x, masked_vis = preprocess_image(image)
-    preds = model_capsnet.predict(x)
-
-    scores = {CLASS_LABELS[i]: float(preds[0][i]) for i in range(len(CLASS_LABELS))}
-
-    # return both: class scores + masked image
-    return scores, masked_vis
+with gr.Blocks() as demo:
+    img_input = gr.Image(type="filepath", label=_change_input_label(None)["label"])
+    out_mask = gr.Image(type="numpy", label="Masked Lung Region")
+    out_label = gr.Label(label="Prediction")
+
+    examples = gr.Examples(
+        label="Example X-rays",
+        examples=img_paths[:], # each element is a list of one string (path):  [[path1], [path2], ... ]
+        example_labels=img_names[:],
+        inputs=img_input,
+        cache_examples=False,
+    )
+    
 
+    # event order of img_input: 
+    # * for uploads: select() → upload() → change() 
+    # * for examples clicks: select() → change()
+    #img_input.select(_change_input_label, inputs=img_input, outputs=img_input) 
+    img_input.change(_change_input_label, inputs=img_input, outputs=img_input) 
+    img_input.change(_clear_outputs, inputs=None, outputs=[out_label, out_mask], queue=False)
+    img_input.change(_predict_and_display, inputs=img_input, outputs=[out_label, out_mask]) 
+    
 
-# ------------------------------------------------------------
-# 4️⃣ Gradio interface
-# ------------------------------------------------------------
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil", label="Upload Chest X-ray"),
-    outputs=[
-        gr.Label(num_top_classes=4, label="Prediction"),
-        gr.Image(type="numpy", label="Masked Lung Region"),
-    ],
-    title="CapsNet 4-Class Lung Disease Classifier",
-    description=(
-        "This demo uses a Capsule Network trained to classify chest X-rays into "
-        "COVID-19, Lung Opacity, Normal, and Viral Pneumonia. "
-        "The masked image shows only the segmented lung region used for classification."
-    ),
-    examples=demo_images
-)
 
 # ------------------------------------------------------------
 # 5️⃣ Launch the app
 # ------------------------------------------------------------
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=True, server_name="127.0.0.1", server_port=7860)

base.py

@@ -0,0 +1,244 @@
+from datasets import load_dataset
+import gradio as gr
+import tensorflow as tf
+from tensorflow.keras.preprocessing.image import img_to_array
+from huggingface_hub import hf_hub_download
+from PIL import Image
+import os
+import numpy as np
+from modelbuilder import capsnet_custom_objects
+from defs import class_to_disease_map
+import pandas as pd
+from typing import Mapping
+from pathlib import Path
+
+
+# ------------------------------------------------------------
+#  Determine the running environment: local machine or huggingface spaces
+# ------------------------------------------------------------
+def running_in_spaces() -> bool:
+    """Return True if app is running inside Hugging Face Spaces."""
+    return (
+        os.environ.get("SPACE_ID") is not None or os.environ.get("SYSTEM") == "spaces"
+    )
+
+is_spaces = running_in_spaces()
+
+if is_spaces:
+    print(f"Running in Hugging Face Spaces environment.")
+else:
+    print(f"Running on local machine environment:{os.environ.get('COMPUTERNAME','Unknown')}")
+    
+# -------CONSTANTS-------#
+TARGET_SIZE = (256, 256)  # target size for masked images
+
+# ------------------------------------------------------------
+# 1️⃣ Load the models from Hugging Face Hub
+# capsnet for disease classification and GAN for lung segmentation/masking
+# ------------------------------------------------------------
+gan_model_path = None
+capsnet_model_path = None
+dataset = None
+data_dir = None
+
+if is_spaces:
+    # huggingface datasets is preinstalled in Spaces
+
+    data_dir = "external_xrays_299x299"
+    dataset = load_dataset(
+        "valste/lung-disease-xrays", data_dir=data_dir, split="train"
+    )
+
+    gan_model_path = hf_hub_download(
+        repo_id="valste/lung-segmentation-gan", filename="model.keras"
+    )
+    capsnet_model_path = hf_hub_download(
+        repo_id="valste/capsnet-4class-lung-disease-classifier", filename="model.keras"
+    )
+
+else:
+    # local machine
+    capsnet_model_path = os.path.join(
+        ".", "models", "capsnet-4class-lung-disease-classifier", "model.keras"
+    )
+    #capsnet_knl5_path = Path(r"C:\Users\User\VS_CODE_WORKSPACES\DataScience\ds_aug24_lung_desease_classification\mlruns_capsnet\923962217935764323\af22305e09b44dd686b2405b0f6c01d5\artifacts\model\data\model.keras") #conv2d_1 krnl=5
+    gan_model_path = os.path.join(".", "models", "lung-segmentation-gan", "model.keras")
+    data_dir = os.path.join(".", "data", "external_xrays_299x299")
+    dataset = load_dataset(
+        "imagefolder", data_dir=data_dir, split="train"  # path to your local folder
+    )
+
+model_gan = tf.keras.models.load_model(gan_model_path, compile=False)
+model_capsnet = tf.keras.models.load_model(
+    capsnet_model_path, 
+    #capsnet_knl5_path,
+    custom_objects=capsnet_custom_objects, compile=False
+)
+
+# ------------------------------------------------------------
+# 2️⃣ Load sample X-ray images from your dataset
+# ------------------------------------------------------------
+imgs=[]
+img_paths = [] 
+img_names = []
+
+class DemoException(Exception):
+    pass
+    
+
+for ex in dataset:
+    if "image" in ex:
+        imgs.append(ex["image"])
+        path = getattr(ex["image"], "filename", None) # string
+        if path:
+            img_paths.append([path])
+            img_names.append(os.path.basename(path))
+        else:
+            raise DemoException("Missing path")
+    else:
+        raise DemoException("Dataset examples do not contain 'image' field.")  
+    
+
+# ------------------------------------------------------------
+# 3️⃣ Define preprocessing and inference function
+# ------------------------------------------------------------
+  
+def create_binary_mask(img: Image.Image, seg_model=model_gan, target_size=TARGET_SIZE) -> np.ndarray:
+    """Create a binary mask from a PIL Image."""
+        # --- 1) Grayscale base image ---
+    img_gray = img.convert("L")
+    #print(f" image size: {img_gray.size}")
+    img_gray = img_gray.resize(target_size, Image.BILINEAR)
+    #print(f" image rescaled size: {img_gray.size}")
+
+    # float32 in [0, 1]
+    gray_array = np.array(img_gray, dtype=np.float32) / 255.0
+    #print(f" gray_array shape: {gray_array.shape}")
+
+    # --- 2) GAN input ---
+    gan_input = gray_array[..., np.newaxis]          # (H, W, 1)
+    gan_input = np.expand_dims(gan_input, axis=0)    # (1, H, W, 1)
+    #print(f" GAN input shape (before prediction): {gan_input.shape}")
+    # --- 3) Run segmentation GAN ---
+    prediction = seg_model.predict(gan_input)
+    mask = (prediction[0, :, :, 0] > 0.5).astype(np.uint8) * 255  # take first batch & channel, threshold at 0.5, get 0/255 mask
+    mask = (mask > 127).astype(np.uint8)  # ensure strictly binary mask with values 0 or 1
+    
+    return mask
+
+def create_masked_img(img: Image.Image, seg_model=model_gan, target_size=TARGET_SIZE) -> tuple[tf.Tensor, np.ndarray, np.ndarray]:
+
+    """Create masked image tensor for CapsNet input from a PIL Image."""
+    img = img.resize(target_size, Image.BILINEAR)
+    mask = create_binary_mask(img, seg_model=seg_model, target_size=target_size)  # (H, W), values in {0, 1}      
+    img_arr = img_to_array(img).astype(np.float32)  # (H,W,3)
+    
+    # --- 4) Apply mask ---
+    masked = img_arr * mask[..., None]     # still float # (H,W,3)
+    masked = masked.astype(np.uint8)       # back to 0–255 image # (H,W,3)
+
+    # --- 5) Prepare input for CapsNet ---
+   # masked_3ch = np.repeat(masked[..., np.newaxis], 3, axis=-1)  # (H,W,3)
+
+    x = tf.convert_to_tensor(
+        np.expand_dims(masked, axis=0),  # (1, H, W, 3)
+        dtype=tf.float32,                        # ✅ float32, not uint8
+    )
+
+    return x, masked, mask
+
+
+
+def to_probabilities(predictions, model_type: str = "capsnet"):
+    """
+    Convert model outputs to class probabilities.
+
+    Parameters
+    ----------
+    dis_confs : np.ndarray
+        Raw model outputs.
+        For CapsNet this is expected to be of shape (n, n_classes, 1)
+        where the last dim is just a singleton and dis_confs need normalization means the sum along class axis is not equal to 1.
+    model_type : str
+        "capsnet" → normalize along class axis and squeeze last dim.
+        anything else → return dis_confs unchanged as np.ndarray.
+
+    Returns
+    -------
+    np.ndarray
+        Probabilities of shape (n, n_classes) for capsnet,
+        or np.asarray(dis_confs) for other models.
+    """
+    predictions = np.asarray(predictions)
+
+    if model_type != "capsnet":
+        # Assume other models (e.g. softmax CNN) already return probs.
+        return predictions
+
+    # --- CapsNet branch ---
+    # Ensure we have a 3D tensor with the singleton last dimension
+    if predictions.ndim == 2:
+        # (n, n_classes) -> (n, n_classes, 1)
+        predictions = predictions[..., np.newaxis]
+    elif predictions.ndim != 3:
+        raise ValueError(f"Expected dis_confs to have 2 or 3 dims for capsnet, got shape {predictions.shape}")
+
+    # Normalize along the class axis (axis=1)
+    sums = np.sum(predictions, axis=1, keepdims=True)
+
+    # Avoid division by zero just in case
+    eps = 1e-12
+    sums = np.where(sums == 0, eps, sums)
+
+    predictions = predictions / sums
+
+    # Remove the last singleton dim: (n, n_classes, 1) -> (n, n_classes)
+    predictions = predictions.squeeze(-1)
+
+    return predictions
+
+
+
+def map_probabilities_to_classes(probabilities, class_map):
+    """
+    Map a probability vector to a dict {class_name: probability}.
+
+    Parameters
+    ----------
+    probabilities : array-like
+        Shape (n_classes,) – probabilities after softmax or normalization.
+    class_map : dict
+        Mapping from class index → class name.
+
+    Returns
+    -------
+    dict
+        {class_name: float(probability)} sorted by descending probability.
+    """
+    probs = np.asarray(probabilities).flatten()
+
+    if len(probs) != len(class_map):
+        raise ValueError(f"Probability length {len(probs)} does not match class_map size {len(class_map)}")
+
+    mapped = {class_map[i]: float(probs[i]) for i in range(len(probs))}
+
+    # Sort by highest probability
+    mapped = dict(sorted(mapped.items(), key=lambda x: x[1], reverse=True))
+
+    return mapped
+
+
+
+def predict(img_path: str) -> tuple[str, np.ndarray, np.ndarray]:
+    img = Image.open(img_path)
+    x, masked_vis, _ = create_masked_img(img)
+    # ---> ok
+    preds = model_capsnet.predict(x, verbose=1) #(1,4,1) aka (batch_size, 4_class_prob, 1)
+    probs = to_probabilities(preds, model_type="capsnet")  # shape (1,4)
+    probs = probs[0]  # shape (4,)
+      
+    prob_dict = map_probabilities_to_classes(probs, class_to_disease_map)
+    
+    filename_out = os.path.basename(img_path)
+
+    return filename_out, masked_vis, prob_dict

modelbuilder.py

@@ -128,6 +128,8 @@ class ModelBuilder:
     def build_capsnet(self, inputs, **params):
         """
         Build a Capsule Network model for four class lung iseases classification: COVID, Normal, Pneumonia and Opacity.
+        The batch dimension is always None internally → full input shape is (None, 256, 256, 1).
+        The output shape is (None, 4, 1) 
         Args:
             name (_type_): _description_
             first_Conv2DKernel_size (int, optional): _description_. Defaults to 10.
@@ -269,7 +271,7 @@ class PrimaryCaps(layers.Layer):
 
 @register_keras_serializable()
 class DigitCaps(layers.Layer):
-    # DigitCaps Layer / Higher-level capsules (e.g. detecting objects like digits or lungs)
+    # DigitCaps Layer / Higher-level capsules (e.g. detecting objects like animals or lungs)
 
     def __init__(self, num_capsule, dim_capsule, routing_iters=3, **kwargs):
         super(DigitCaps, self).__init__(**kwargs)