Views
No views yet
1import torch
2from miccai_brats.models.unet.unet import PicoUNet
3
4# 1. Instantiate Model
5model = PicoUNet(in_channels=2, num_classes=4)
6
7# 2. Load Checkpoint
8checkpoint = torch.load("best_model.pth", map_location="cpu")
9model.load_state_dict(checkpoint['model_state_dict'])
10
11model.eval()
12
13# 3. Inference
14input_tensor = torch.randn(1, 2, 128, 128) # Fake Batch of [FLAIR, T1Ce]
15with torch.no_grad():
16 output = model(input_tensor)
17 prediction = torch.argmax(output, dim=1)
18
19print("Predicted shape:", prediction.shape) # [1, 128, 128]