Spaces:
Sleeping
Sleeping
File size: 3,350 Bytes
906fcb9 caf6ee7 906fcb9 1baebae 906fcb9 1baebae caf6ee7 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae caf6ee7 1baebae caf6ee7 906fcb9 efc95db a1cc9d3 f128aa8 1d11192 a1cc9d3 efc95db 906fcb9 efc95db 906fcb9 a1cc9d3 906fcb9 efc95db a1cc9d3 906fcb9 | 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 | from __future__ import annotations
import torch
import torch.nn as nn
from monai.utils.module import optional_import
models, _ = optional_import("torchvision.models")
class SimpleNN(nn.Module):
"""
A simple Multi-Layer Perceptron (MLP) for binary classification.
This network consists of two hidden layers with ReLU activation and a dropout layer,
followed by a final sigmoid activation for probability output.
Args:
input_dim (int): The number of input features.
"""
def __init__(self, input_dim: int) -> None:
super().__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, 256),
nn.ReLU(),
nn.Linear(256, 128),
nn.ReLU(),
nn.Dropout(p=0.3),
nn.Linear(128, 1),
)
def forward(self, x):
"""
Forward pass of the classifier.
Args:
x (torch.Tensor): Input tensor of shape (Batch, input_dim).
Returns:
torch.Tensor: Output probabilities of shape (Batch, 1).
"""
return self.net(x)
class CSPCAModel(nn.Module):
"""
Clinically Significant Prostate Cancer (csPCa) risk prediction model using a MIL backbone.
This model repurposes a pre-trained Multiple Instance Learning (MIL) backbone (originally
designed for PI-RADS prediction) for binary csPCa risk assessment. It utilizes the
backbone's feature extractor, transformer, and attention mechanism to aggregate instance-level
features into a bag-level embedding.
The original fully connected classification head of the backbone is replaced by a
custom :class:`SimpleNN` head for the new task.
Args:
backbone (nn.Module): A pre-trained MIL model. The backbone must possess the
following attributes/sub-modules:
- ``net``: The CNN feature extractor.
- ``transformer``: A sequence modeling module.
- ``attention``: An attention mechanism for pooling.
- ``myfc``: The original fully connected layer (used to determine feature dimensions).
Attributes:
fc_cspca (SimpleNN): The new classification head for csPCa prediction.
backbone: The MIL based PI-RADS classifier.
"""
def __init__(self, backbone: nn.Module) -> None:
super().__init__()
self.backbone = backbone
self.clinical_dim = 2
self.projection_dim = 32
self.clinical_projection = nn.Sequential(
nn.Linear(self.clinical_dim, self.projection_dim),
nn.ReLU(),
nn.LayerNorm(self.projection_dim),
# nn.BatchNorm1d(self.projection_dim), # Helps stabilize the merged scale
)
self.fc_dim = backbone.myfc.in_features
self.fc_cspca = SimpleNN(input_dim=self.fc_dim + self.projection_dim)
def forward(self, x, psa_data):
sh = x.shape
x = x.reshape(sh[0] * sh[1], sh[2], sh[3], sh[4], sh[5])
x = self.backbone.net(x)
x = x.reshape(sh[0], sh[1], -1)
x = self.backbone.transformer(x)
a = self.backbone.attention(x)
a = torch.softmax(a, dim=1)
x = torch.sum(x * a, dim=1)
psa_features = self.clinical_projection(psa_data)
x = torch.cat((x, psa_features), dim=1)
x = self.fc_cspca(x)
return x
|