Spaces:
Runtime error
Runtime error
| """Evaluate a DeepLab semantic segmentation checkpoint on a YOLO-polygon dataset.""" | |
| from __future__ import annotations | |
| import argparse | |
| import csv | |
| import json | |
| import os | |
| import random | |
| import sys | |
| from pathlib import Path | |
| import torch | |
| from torch.utils.data import DataLoader | |
| PROJECT_ROOT = Path(__file__).resolve().parents[1] | |
| SCRIPT_DIR = Path(__file__).resolve().parent | |
| if str(SCRIPT_DIR) not in sys.path: | |
| sys.path.insert(0, str(SCRIPT_DIR)) | |
| os.environ.setdefault("TORCH_HOME", str(PROJECT_ROOT / ".torch")) | |
| from train_deeplab_semseg import YoloSemanticDataset, build_model, evaluate, load_compatible_state_dict | |
| from segmentation_utils import load_dataset_class_names | |
| DEFAULT_DATASET = PROJECT_ROOT / "data" / "processed" / "foodseg103_target_yolo" | |
| DEFAULT_CHECKPOINT = ( | |
| PROJECT_ROOT | |
| / "runs" | |
| / "foodseg103_target" | |
| / "deeplabv3_r50_target_weightprobe2_i384" | |
| / "best.pt" | |
| ) | |
| def parse_args() -> argparse.Namespace: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT) | |
| parser.add_argument("--dataset", type=Path, default=DEFAULT_DATASET) | |
| parser.add_argument("--split", choices=["train", "val"], default="val") | |
| parser.add_argument("--imgsz", type=int, default=384) | |
| parser.add_argument("--batch", type=int, default=8) | |
| parser.add_argument("--workers", type=int, default=0) | |
| parser.add_argument("--limit", type=int) | |
| parser.add_argument("--sample-seed", type=int) | |
| parser.add_argument("--backbone", choices=["resnet50", "resnet101"], default="resnet50") | |
| parser.add_argument("--stems-file", type=Path) | |
| parser.add_argument("--output", type=Path) | |
| return parser.parse_args() | |
| def main() -> None: | |
| args = parse_args() | |
| class_names = load_dataset_class_names(args.dataset) | |
| num_classes_with_bg = len(class_names) + 1 | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| dataset = YoloSemanticDataset(args.dataset, args.split, args.imgsz, None, train=False) | |
| if args.stems_file: | |
| with args.stems_file.open(newline="", encoding="utf-8") as f: | |
| stems = [row["stem"] for row in csv.DictReader(f)] | |
| dataset.image_paths = [args.dataset / "images" / args.split / f"{stem}.jpg" for stem in stems] | |
| if args.sample_seed is not None: | |
| rng = random.Random(args.sample_seed) | |
| rng.shuffle(dataset.image_paths) | |
| if args.limit is not None: | |
| dataset.image_paths = dataset.image_paths[: args.limit] | |
| loader = DataLoader( | |
| dataset, | |
| batch_size=args.batch, | |
| shuffle=False, | |
| num_workers=args.workers, | |
| pin_memory=torch.cuda.is_available(), | |
| ) | |
| model = build_model(num_classes_with_bg, pretrained=False, backbone=args.backbone).to(device) | |
| load_compatible_state_dict(model, args.checkpoint, device) | |
| miou, ious = evaluate(model, loader, device, num_classes_with_bg, class_names) | |
| print(f"split={args.split}") | |
| print(f"limit={args.limit if args.limit is not None else 'all'}") | |
| for name, iou in ious.items(): | |
| print(f"{name:12s} IoU={iou:.4f}") | |
| print(f"core mIoU={miou:.4f}") | |
| if args.output: | |
| payload = { | |
| "split": args.split, | |
| "limit": args.limit if args.limit is not None else "all", | |
| "imgsz": args.imgsz, | |
| "checkpoint": str(args.checkpoint), | |
| "stems_file": str(args.stems_file) if args.stems_file else None, | |
| "ious": ious, | |
| "core_miou": miou, | |
| } | |
| args.output.parent.mkdir(parents=True, exist_ok=True) | |
| args.output.write_text(json.dumps(payload, indent=2), encoding="utf-8") | |
| print(f"Wrote evaluation result: {args.output}") | |
| if __name__ == "__main__": | |
| main() | |