Views
No views yet
1import torch
2from model import AttentionUNet
3
4model = AttentionUNet(img_ch=1, output_ch=4)
5state_dict = torch.load("fold_1_model.pth", map_location="cpu", weights_only=False)
6if 'model_state_dict' in state_dict:
7 state_dict = state_dict['model_state_dict']
8model.load_state_dict(state_dict)
9model.eval()
10
11# Input: [batch, 1, 256, 256] normalized to mean=0.5, std=0.5
12img_tensor = torch.randn(1, 1, 256, 256)
13with torch.no_grad():
14 output = model(img_tensor) # [batch, 4, 256, 256]
15 pred = torch.argmax(output, dim=1) # [batch, 256, 256]| File | Description |
|---|---|
model.py | Model architecture (AttentionUNet) |
fold_1_model.pth - fold_5_model.pth | Trained weights for each CV fold |