Spaces:
Runtime error
Runtime error
File size: 2,751 Bytes
1cf3825 | 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 | """Evaluate a saved Attention U-Net checkpoint on val + test splits.
Re-runs the evaluation loop from train_segmentation_torch.py against the
checkpoint at segmentation_artifacts/attention_unet/best_model.pt and writes
evaluation_metrics.json. Useful after training crashes / is interrupted.
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
import torch
from torch.utils.data import DataLoader
_REPO_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(_REPO_ROOT))
from src.segmentation_torch import AttentionUNet # noqa: E402
from src.train_segmentation_torch import SegDataset, _evaluate # noqa: E402
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--weights', default='segmentation_artifacts/attention_unet/best_model.pt')
parser.add_argument('--data_dir', default='dataset_real')
parser.add_argument('--batch_size', type=int, default=8)
parser.add_argument('--threshold', type=float, default=0.5)
parser.add_argument('--device', default='cuda')
args = parser.parse_args()
if args.device == 'cuda' and not torch.cuda.is_available():
args.device = 'cpu'
device = torch.device(args.device)
print('Using device:', device)
ckpt = torch.load(args.weights, map_location=device, weights_only=False)
cfg = ckpt.get('config', {}) or {}
image_size = int(cfg.get('image_size', 256))
print(f'Loaded checkpoint from epoch {ckpt.get("epoch")} (image_size={image_size}, '
f'base_filters={cfg.get("base_filters", 32)})')
model = AttentionUNet(
in_channels=3,
base_filters=int(cfg.get('base_filters', 32)),
dropout=float(cfg.get('dropout', 0.2)),
).to(device)
model.load_state_dict(ckpt['state_dict'])
model.eval()
data_dir = Path(args.data_dir)
payload = {'source_checkpoint_epoch': ckpt.get('epoch'), 'config': cfg}
for split in ['val', 'test']:
split_dir = data_dir / split
if not split_dir.exists():
continue
ds = SegDataset(split_dir, image_size, augment=False)
loader = DataLoader(ds, batch_size=args.batch_size, shuffle=False, num_workers=0)
metrics = _evaluate(model, loader, device, threshold=args.threshold)
payload[split] = metrics
print(f'\n{split}: {len(ds)} samples')
for k, v in metrics.items():
print(f' {k}: {v:.4f}' if isinstance(v, float) else f' {k}: {v}')
out = Path(args.weights).parent / 'evaluation_metrics.json'
out.write_text(json.dumps(payload, indent=2), encoding='utf-8')
print(f'\nWrote {out}')
if __name__ == '__main__':
main()
|