update app.py

app.py

@@ -33,22 +33,263 @@ import albumentations as A
 
 from huggingface_hub import login, hf_hub_download
 
-from unet_models import UNet
-from segformer_model import SegFormer
-from multimodal_model import MultimodalModel
-
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
+class DoubleConv(nn.Module):
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True))
+    def forward(self, x):
+        return self.double_conv(x)
+    
+class Down(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv(in_channels, out_channels))
+    def forward(self, x):
+        return self.maxpool_conv(x)
+    
+class Up(nn.Module):
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, in_channels//2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels//2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels)
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+        
+        x1 = F.pad(x1, [diffX//2, diffX-diffX//2,
+                        diffY//2, diffY-diffY//2])
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+class OutConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(OutConv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=1),
+            nn.Sigmoid())
+    def forward(self, x):
+        return self.conv(x)
+
+class UNet(nn.Module):
+    def __init__(self, n_channels, n_classes, bilinear=True):
+        super(UNet, self).__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+        self.bilinear = bilinear
+        
+        self.inc = DoubleConv(n_channels, 64)
+        self.down1 = Down(64, 128)
+        self.down2 = Down(128, 256)
+        self.down3 = Down(256, 512)
+        factor = 2 if bilinear else 1
+        self.down4 = Down(512, 1024//factor)
+        self.up1 = Up(1024, 512//factor, bilinear)
+        self.up2 = Up(512, 256//factor, bilinear)        
+        self.up3 = Up(256, 128//factor, bilinear)        
+        self.up4 = Up(128, 64, bilinear)        
+        self.outc = OutConv(64, n_classes)
+    
+    def forward(self, x):
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x = self.up1(x5, x4)
+        x = self.up2(x, x3)
+        x = self.up3(x, x2)
+        x = self.up4(x, x1)
+        logits = self.outc(x)
+        return logits
+
 model = UNet(3, 1).to(DEVICE)
 out = model(torch.randn(1, 3, 128, 128).to(DEVICE))
 print(out.shape)
 
+
+class SqueezeExcitation(nn.Module):
+    def __init__(self, channels: int, reduction: int = 16):
+        super(SqueezeExcitation, self).__init__()
+        self.se = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(channels, channels // reduction, kernel_size=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channels // reduction, channels, kernel_size=1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x * self.se(x)
+
+class SegFormer(nn.Module):
+    def __init__(self, n_channels: int, n_classes: int, pretrained_model: str = "nvidia/mit-b5"):
+        super(SegFormer, self).__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+        
+        config = SegformerConfig.from_pretrained(
+            pretrained_model,
+            num_channels=n_channels,
+            num_labels=n_classes,
+            hidden_dropout_prob=0.3, 
+            attention_probs_dropout_prob=0.3,  
+            drop_path_rate=0.1  
+        )
+
+        self.segformer = SegformerForSemanticSegmentation.from_pretrained(
+            pretrained_model,
+            config=config,
+            ignore_mismatched_sizes=True
+        )
+        
+        if n_channels != 3:
+            self.segformer.segformer.encoder.patch_embeddings[0].proj = nn.Conv2d(
+                n_channels, config.hidden_sizes[0], kernel_size=7, stride=4, padding=2
+            )
+
+        self.segformer.decode_head.classifier = nn.Sequential(
+            nn.Conv2d(config.decoder_hidden_size, config.decoder_hidden_size // 2, kernel_size=3, padding=1),
+            nn.BatchNorm2d(config.decoder_hidden_size // 2, momentum=0.05), 
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(0.4),
+            nn.Conv2d(config.decoder_hidden_size // 2, n_classes, kernel_size=1)
+        )
+
+        self.fpn = nn.ModuleList([
+            nn.Sequential(
+                nn.Conv2d(h, 128, kernel_size=1),
+                nn.BatchNorm2d(128, momentum=0.05),  
+                nn.ReLU(inplace=True),
+                SqueezeExcitation(128)
+            ) for h in config.hidden_sizes
+        ])
+
+        self.fusion = nn.Sequential(
+            nn.Conv2d(128 * len(config.hidden_sizes), 256, kernel_size=3, padding=1),
+            nn.BatchNorm2d(256, momentum=0.05), 
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(0.3), 
+            nn.Conv2d(256, 256, kernel_size=3, padding=1),
+            nn.BatchNorm2d(256, momentum=0.05),
+            nn.ReLU(inplace=True)
+        )
+        
+        self.fusion_residual = nn.Conv2d(128 * len(config.hidden_sizes), 256, kernel_size=1)
+
+        self.refinement = nn.Sequential(
+            nn.Conv2d(256 + n_classes, 128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128, momentum=0.05),
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(0.2),  
+            nn.Conv2d(128, n_classes, kernel_size=1)
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        input_size = x.size()[2:]
+
+        outputs = self.segformer(pixel_values=x)
+        logits = outputs.logits
+
+        encoder_outputs = self.segformer.segformer.encoder(pixel_values=x, output_hidden_states=True)
+        hidden_states = encoder_outputs.hidden_states
+
+        fpn_feats = []
+        for i, (feat, layer) in enumerate(zip(hidden_states, self.fpn)):
+            f = layer(feat)
+            f = F.interpolate(f, size=logits.shape[2:], mode="bilinear", align_corners=False)
+            fpn_feats.append(f)
+
+        fused = torch.cat(fpn_feats, dim=1)
+        residual = self.fusion_residual(fused)
+        fused = self.fusion(fused)
+        fused = fused + residual  
+
+        logits = F.interpolate(logits, size=input_size, mode="bilinear", align_corners=False)
+        fused = F.interpolate(fused, size=input_size, mode="bilinear", align_corners=False)
+
+        concat = torch.cat([fused, logits], dim=1)
+        out = self.refinement(concat)
+
+        return out  
+
 model = SegFormer(n_channels=3, n_classes=1).to(DEVICE)
 out = model(torch.randn(1, 3, 128, 128).to(DEVICE))
 print(out.shape)
 
+
+class MultimodalModel(nn.Module):
+    def __init__(self, num_numeric_features, num_classes):
+        super(MultimodalModel, self).__init__()
+        
+        self.vit = models.vit_b_16(pretrained=True)
+        self.vit.heads = nn.Identity()
+        
+        self.swin_b = models.swin_b(pretrained=True)
+        self.swin_b.head = nn.Identity()
+        
+        self.swinv2_b = models.swin_v2_b(pretrained=True)
+        self.swinv2_b.head = nn.Identity()
+        
+        self.numeric_branch = nn.Sequential(
+            nn.Conv1d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm1d(16),
+            nn.ReLU(),
+            nn.MaxPool1d(kernel_size=2, stride=2),
+            nn.Conv1d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm1d(32),
+            nn.ReLU(),
+            nn.MaxPool1d(kernel_size=2, stride=2),
+            nn.Flatten(),
+            nn.Linear((num_numeric_features // 4) * 32, 64),
+            nn.BatchNorm1d(64),
+            nn.ReLU(),
+            nn.Linear(64, num_classes)
+        )
+        
+        self.image_fc = nn.Sequential(
+            nn.Linear(768 + 1024 + 1024, 128),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(128, num_classes)
+        )
+    
+    def forward(self, image, numeric_data):
+        vit_features = self.vit(image)
+        
+        swin_b_features = self.swin_b(image)
+        
+        swinv2_b_features = self.swinv2_b(image)
+        
+        combined_image_features = torch.cat((vit_features, swin_b_features, swinv2_b_features), dim=1)  # Shape: (N, 2816)
+        combined_image_output = self.image_fc(combined_image_features)
+        
+        numeric_data = numeric_data.unsqueeze(1)
+        numeric_output = self.numeric_branch(numeric_data)
+        
+        final_output = 0.95 * combined_image_output + 0.05 * numeric_output
+        
+        return final_output
+
+
 secret_value_0 = os.getenv("carc_hf_token") 
 login(token=secret_value_0)
 print("Logged in successfully")
@@ -102,7 +343,6 @@ segformer_macula_segmentation_model.eval()
 segformer_macula_segmentation_model = segformer_macula_segmentation_model.to(DEVICE)
 print("Loaded SegFormer macula segmentation model")
 
-
 multimodal_glaucoma_classification_model = MultimodalModel(num_numeric_features=23, num_classes=2)
 multimodal_glaucoma_classification_model_path = hf_hub_download(
     repo_id="rprkh/multimodal_glaucoma_classification",
@@ -191,7 +431,7 @@ def extract_macula_area(image_pil):
 
     with torch.no_grad():
         output = segformer_macula_segmentation_model(img_tensor)
-        mask = (output > 0.5).float()
+        mask = (output > 0.5).float()   # (1,1,H,W)
 
     macula_area = mask.sum().item()
     return macula_area
@@ -243,10 +483,7 @@ def compute_vasculature_density(image_pil, model, device, threshold=0.05, radius
         roi_area    = roi_tensor.sum().item()
         density = vessel_area / roi_area if roi_area > 0 else 0.0
 
-        overlay_image = blend_image_with_mask(image_pil, masked)  # Green vessels
-        vasculature_density_masked = Image.fromarray((overlay_image * 255).astype(np.uint8))
-
-        return density, vasculature_density_masked
+        return density
 
 def create_circular_roi_mask(image_shape, radius_ratio=0.95):
     h, w = image_shape
@@ -405,7 +642,7 @@ def predict_all_diameters(image_path):
 
     vis_rgb = cv2.cvtColor(vis_bgr, cv2.COLOR_BGR2RGB) / 255.0
 
-    vd, vasculature_density_masked = compute_vasculature_density(
+    vd = compute_vasculature_density(
             image_pil=image_pil,
             model=retinal_vasculature_segmentation_model,
             device=DEVICE,
@@ -532,65 +769,65 @@ def predict_all_diameters(image_path):
         return measurement
 
     result_text = f"""
-        <div id="results_container">
-        <div id="results_table" style="display:flex; gap:7px;">
-        <div style="flex:1;">
-        <table>
-        <tr><th>Measurement</th><th>Value</th></tr>
-
-        <tr><td><b>Retina Diameter</b></td><td>{convert_to_mm(retina_diameter)} mm</td></tr>
-        <tr><td><b>Optic Cup Diameter</b></td><td>{convert_to_mm(cup_diameter)} mm</td></tr>
-        <tr><td><b>Optic Disc Diameter</b></td><td>{convert_to_mm(disc_diameter)} mm</td></tr>
-        <tr><td><b>Macular Diameter</b></td><td>{convert_to_mm(macula_diameter)} mm</td></tr>
-        <tr><td><b>Vasculature Density</b></td><td>{round(vd * 100, 3)}%</td></tr>
-
-        <tr><td><b>Retina Radius</b></td><td>{convert_to_mm(retina_radius)} mm</td></tr>
-        <tr><td><b>Retina Area</b></td><td>{convert_to_mm2(retina_area)} mm<sup>2</sup></td></tr>
-        <tr><td><b>Retina Circumference</b></td><td>{convert_to_mm(retina_circumference)} mm</td></tr>
-        </table>
-        </div>
-
-        <div style="flex:1;">
-        <table>
-        <tr><th>Measurement</th><th>Value</th></tr>
-
-        <tr><td><b>Optic Cup Radius</b></td><td>{convert_to_mm(optic_cup_radius)} mm</td></tr>
-        <tr><td><b>Optic Cup Area</b></td><td>{convert_to_mm2(optic_disc_area)} mm<sup>2</sup></td></tr>
-        <tr><td><b>Optic Cup Circumference</b></td><td>{convert_to_mm(optic_cup_circumference)} mm</td></tr>
-
-        <tr><td><b>Optic Disc Radius</b></td><td>{convert_to_mm(optic_disc_radius)} mm</td></tr>
-        <tr><td><b>Optic Disc Area</b></td><td>{convert_to_mm2(optic_cup_area)} mm<sup>2</sup></td></tr>
-        <tr><td><b>Optic Disc Circumference</b></td><td>{convert_to_mm(optic_disc_circumference)} mm</td></tr>
-
-        <tr><td><b>Macula Radius</b></td><td>{convert_to_mm(macula_radius)} mm</td></tr>
-        <tr><td><b>Macula Area</b></td><td>{convert_to_mm(macula_area)} mm</td></tr>
-
-        </table>
-        </div>
-
-        <div style="flex:1;">
-        <table>
-        <tr><th>Measurement</th><th>Value</th></tr>
-
-        <tr><td><b>Macula Circumference</b></td><td>{convert_to_mm(macula_circumference)} mm</td></tr>
-
-        <tr><td><b>Optic Disc to Retina Diameter Ratio</b></td><td>{round_measurement_to_3_dp(optic_disc_to_retina_diameter_ratio)}</td></tr>
-        <tr><td><b>Optic Disc to Retina Area Ratio</b></td><td>{round_measurement_to_3_dp(optic_disc_to_retina_area_ratio)}</td></tr>
-
-        <tr><td><b>Optic Cup to Disc Diameter Ratio</b></td><td>{round_measurement_to_3_dp(optic_cup_to_disc_diameter_ratio)}</td></tr>
-        <tr><td><b>Optic Cup to Disc Area Ratio</b></td><td>{round_measurement_to_3_dp(optic_cup_to_disc_area_ratio)}</td></tr>
-        <tr><td><b>Optic Cup to Retina Diameter Ratio</b></td><td>{round_measurement_to_3_dp(optic_cup_to_retina_diameter_ratio)}</td></tr>
-        <tr><td><b>Optic Cup to Retina Area Ratio</b></td><td>{round_measurement_to_3_dp(optic_cup_to_retina_area_ratio)}</td></tr>
-        </table>
-        </div>
-        </div>
-
-        <h3>Predicted Class: {prediction}</h3>
-        <h3>Confidence: {round(confidence * 100, 3)}%</h3>
-        <div>
-    """
+<div id="results_container">
+<div id="results_table" style="display:flex; gap:7px;">
+<div style="flex:1;">
+<table>
+<tr><th>Measurement</th><th>Value</th></tr>
+
+<tr><td><b>Retina Diameter</b></td><td>{convert_to_mm(retina_diameter)} mm</td></tr>
+<tr><td><b>Optic Cup Diameter</b></td><td>{convert_to_mm(cup_diameter)} mm</td></tr>
+<tr><td><b>Optic Disc Diameter</b></td><td>{convert_to_mm(disc_diameter)} mm</td></tr>
+<tr><td><b>Macular Diameter</b></td><td>{convert_to_mm(macula_diameter)} mm</td></tr>
+<tr><td><b>Vasculature Density</b></td><td>{round(vd * 100, 3)}%</td></tr>
+
+<tr><td><b>Retina Radius</b></td><td>{convert_to_mm(retina_radius)} mm</td></tr>
+<tr><td><b>Retina Area</b></td><td>{convert_to_mm2(retina_area)} mm<sup>2</sup></td></tr>
+<tr><td><b>Retina Circumference</b></td><td>{convert_to_mm(retina_circumference)} mm</td></tr>
+</table>
+</div>
+
+<div style="flex:1;">
+<table>
+<tr><th>Measurement</th><th>Value</th></tr>
+
+<tr><td><b>Optic Cup Radius</b></td><td>{convert_to_mm(optic_cup_radius)} mm</td></tr>
+<tr><td><b>Optic Cup Area</b></td><td>{convert_to_mm2(optic_disc_area)} mm<sup>2</sup></td></tr>
+<tr><td><b>Optic Cup Circumference</b></td><td>{convert_to_mm(optic_cup_circumference)} mm</td></tr>
+
+<tr><td><b>Optic Disc Radius</b></td><td>{convert_to_mm(optic_disc_radius)} mm</td></tr>
+<tr><td><b>Optic Disc Area</b></td><td>{convert_to_mm2(optic_cup_area)} mm<sup>2</sup></td></tr>
+<tr><td><b>Optic Disc Circumference</b></td><td>{convert_to_mm(optic_disc_circumference)} mm</td></tr>
+
+<tr><td><b>Macula Radius</b></td><td>{convert_to_mm(macula_radius)} mm</td></tr>
+<tr><td><b>Macula Area</b></td><td>{convert_to_mm(macula_area)} mm</td></tr>
+
+</table>
+</div>
+
+<div style="flex:1;">
+<table>
+<tr><th>Measurement</th><th>Value</th></tr>
+
+<tr><td><b>Macula Circumference</b></td><td>{convert_to_mm(macula_circumference)} mm</td></tr>
+
+<tr><td><b>Optic Disc to Retina Diameter Ratio</b></td><td>{round_measurement_to_3_dp(optic_disc_to_retina_diameter_ratio)}</td></tr>
+<tr><td><b>Optic Disc to Retina Area Ratio</b></td><td>{round_measurement_to_3_dp(optic_disc_to_retina_area_ratio)}</td></tr>
+
+<tr><td><b>Optic Cup to Disc Diameter Ratio</b></td><td>{round_measurement_to_3_dp(optic_cup_to_disc_diameter_ratio)}</td></tr>
+<tr><td><b>Optic Cup to Disc Area Ratio</b></td><td>{round_measurement_to_3_dp(optic_cup_to_disc_area_ratio)}</td></tr>
+<tr><td><b>Optic Cup to Retina Diameter Ratio</b></td><td>{round_measurement_to_3_dp(optic_cup_to_retina_diameter_ratio)}</td></tr>
+<tr><td><b>Optic Cup to Retina Area Ratio</b></td><td>{round_measurement_to_3_dp(optic_cup_to_retina_area_ratio)}</td></tr>
+</table>
+</div>
+</div>
+
+<h3>Predicted Class: {prediction}</h3>
+<h3>Confidence: {round(confidence * 100, 3)}%</h3>
+<div>
+"""
 
-    return result_text, vasculature_density_masked
+    return result_text
 
 
 custom_css = """
@@ -661,17 +898,14 @@ with gr.Blocks(title="Glaucoma Predictor", css=custom_css) as demo:
             image_input = gr.Image(type="filepath", label="Upload Fundus Image", elem_id="image_box")
             btn = gr.Button("Analyze Image", variant="primary", elem_id="prediction_button")
             result_md = gr.Markdown(elem_id="results_container")
-    
-    with gr.Row():
-        vasculature_density_masked = gr.Image(label="Segmentation Visualization", elem_id="Retina Diameter")
 
     btn.click(
         fn=lambda: ("Analyzing... Please wait.",),
-        outputs=[result_md, vasculature_density_masked]
+        outputs=[result_md]
     ).then(
         fn=predict_all_diameters,
         inputs=image_input,
-        outputs=[result_md, vasculature_density_masked]
+        outputs=[result_md]
     )
 
 if __name__ == "__main__":