Hugging Face's logo

valste
/
capsnet-4class-lung-disease-classifier

Keras
Model card Files
xet
Community
1
capsnet-4class-lung-disease-classifier
File size: 5,080 Bytes
f550944
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
 
import gradio as gr
from datasets import load_dataset
import tensorflow as tf
from huggingface_hub import hf_hub_download
from PIL import Image
import os
import numpy as np
from tensorflow.keras.preprocessing.image import img_to_array
from modelbuilder import capsnet_custom_objects


# 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 in local machine environment:{os.environ.get('SYSTEM')}")

# -------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
# ------------------------------------------------------------
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:
    raise NotImplementedError("clone required models locally and adjust paths in here first! ... and remove this line afterwards ;)")
    
    # local machine
    capsnet_model_path = os.path.join(
        ".", "models", "capsnet-4class-lung-disease-classifier", "model.keras"
    )
    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, 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_masked_img(img: Image.Image) -> tuple[np.ndarray, np.ndarray]:
    
    # --- 1) Make a grayscale base image for segmentation ---
    img_gray = img.convert("L")  # grayscale (299, 299)
    img_gray = img_gray.resize(TARGET_SIZE, Image.BILINEAR)  # (256, 256)
    gray_array = np.array(img_gray, dtype=np.float32) / 255.0  # (H, W) in [0,1]

    # --- 2) Build a 1-channel version for GAN input ---
    gan_input = gray_array[..., np.newaxis]  # adding channel dim (H, W) → (H, W, 1)
    gan_input = np.expand_dims(gan_input, axis=0)  # adding batch dim (1, H, W, 1)
   
    # --- 3) Run segmentation GAN to get lung mask ---
    prediction = model_gan.predict(gan_input)  # (1, 256, 256, 1)
    lung_prob = np.squeeze(prediction)  # (H, W)
    mask = (lung_prob > 0.5).astype(np.float32)  # (H, W), 0/1

    # --- 4) Apply mask  ---
    masked_gray = gray_array * mask  # (H, W)
    
    # --- 5) Prepare input for CapsNet: 3 channel + batch dimension---
    masked_gray_ch = np.repeat(
        masked_gray[..., np.newaxis], 3, axis=-1
    )  # channel dim (H, W, 3)
    
    # adding 
    x = np.expand_dims(masked_gray_ch, axis=0).astype(
        np.float32
    )  # (1, H, W, 1) ✅ for CapsNet

    return x, masked_gray


def predict(img_path: str) -> tuple[str, np.ndarray, np.ndarray]:
    
    img = Image.open(img_path)
    x, masked_vis = create_masked_img(img)
    preds = model_capsnet.predict(x, verbose=0)

    preds = np.asarray(preds)
    if preds.ndim > 2:
        preds = np.squeeze(preds, axis=-1) # (1, 4, 1) → (1, 4)
    preds = np.squeeze(preds)  # (4,)

    scores = {CLASS_LABELS[i]: float(preds[i]) for i in range(len(CLASS_LABELS))}
    
    filename_out = os.path.basename(img_path)

    return filename_out, masked_vis, scores