File size: 3,210 Bytes
4365ce6
 
cf26fcf
4365ce6
 
12eeee4
c036ac6
4365ce6
 
 
 
 
3875845
4365ce6
3875845
4365ce6
 
 
 
3875845
4365ce6
3875845
4365ce6
 
 
 
 
 
 
 
c036ac6
4365ce6
 
 
 
 
c3ea62f
ae8ffe4
c2691cf
c5703b8
 
3875845
 
c5703b8
535ec48
c2691cf
c5703b8
3875845
c2691cf
 
 
c5703b8
c3ea62f
 
 
 
 
 
9a825a5
c036ac6
c2691cf
c3ea62f
 
c5703b8
c036ac6
 
c5703b8
 
 
 
564d05c
c5703b8
6502e81
c5703b8
4365ce6
c036ac6
4365ce6
 
 
 
 
c036ac6
4365ce6
 
 
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
import gradio as gr
import torch
import torch.nn.functional as F
import nibabel as nib
import numpy as np

# Define the ConvAutoencoder model structure
class ConvAutoencoder(torch.nn.Module):
    def __init__(self):
        super(ConvAutoencoder, self).__init__()
        # Encoder
        self.encoder = torch.nn.Sequential(
            torch.nn.Conv2d(1, 32, 3, stride=2, padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(32, 64, 3, stride=2, padding=1),
            torch.nn.ReLU()
        )
        # Decoder
        self.decoder = torch.nn.Sequential(
            torch.nn.ConvTranspose2d(64, 32, 2, stride=2, padding=1),
            torch.nn.ReLU(),
            torch.nn.ConvTranspose2d(32, 1, 2, stride=2, padding=1),
            torch.nn.Sigmoid()
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

# Load the pre-trained model
model_path = "conv_autoencoder_model.pth"
model = ConvAutoencoder()
model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")), strict=False)
model.eval()

# Prediction function
def predict(modalities):
    try:
        slices = []
        for modality in modalities:
            nii_data = nib.load(modality).get_fdata()
            slices.append(nii_data)
        
        # Average modalities to create a single-channel image
        slice_avg = np.mean(slices, axis=0)
        
        # Select the slice with maximum non-zero pixels
        non_zero_counts = [np.count_nonzero(slice_avg.take(i, axis=2)) for i in range(slice_avg.shape[2])]
        max_index = np.argmax(non_zero_counts)
        best_slice = slice_avg.take(max_index, axis=2)
        
        # Convert best_slice to a tensor and resize to (256, 256) using interpolation
        tensor_slice = torch.tensor(best_slice).unsqueeze(0).unsqueeze(0).float()  # Shape: [1, 1, H, W]
        
        # Force resize to (256, 256)
        tensor_slice_resized = F.interpolate(tensor_slice, size=(256, 256), mode="bilinear", align_corners=False)
        print(f"Tensor shape after forced resize: {tensor_slice_resized.shape}")  # Debugging
        
        # Run model inference
        with torch.no_grad():
            reconstruction = model(tensor_slice_resized)
        reconstruction_error = torch.abs(reconstruction - tensor_slice_resized).squeeze().numpy()
        
        # Anomaly detection
        error_threshold = 0.1
        anomaly_detected = np.mean(reconstruction_error) > error_threshold
        anomaly_message = "Anomaly Detected" if anomaly_detected else "No Anomaly Detected"
        
        return reconstruction_error, anomaly_message

    except Exception as e:
        print(f"Processing error: {e}")
        return np.zeros((256, 256)), f"Processing error: {e}"

# Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Files(type="filepath", label="Upload NIfTI Modalities (e.g., T1, T2, FLAIR)"),
    outputs=[gr.Image(type="numpy", label="Reconstruction Error"), gr.Textbox(label="Anomaly Detection")],
    title="Brain MRI Anomaly Detection - ConvAutoencoder",
    description="Upload NIfTI files for brain anomaly detection using ConvAutoencoder.",
)

iface.launch()