| """ |
| Evaluate a SAM ViT-H checkpoint (original or fine-tuned) on facade segmentation. |
| """ |
| import os |
| import argparse |
| import json |
| import numpy as np |
| import torch |
| from torch.utils.data import DataLoader |
| from transformers import SamModel, SamProcessor |
| from tqdm import tqdm |
|
|
| from dataset import FacadeDataset, collate_fn |
| from metrics import compute_all_metrics |
|
|
|
|
| @torch.no_grad() |
| def evaluate_checkpoint(checkpoint_path, data_dir, split="test", batch_size=1, output_dir="outputs/eval"): |
| os.makedirs(output_dir, exist_ok=True) |
|
|
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| print(f"Device: {device}") |
|
|
| processor = SamProcessor.from_pretrained("facebook/sam-vit-huge") |
| model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device) |
|
|
| if checkpoint_path and os.path.exists(checkpoint_path): |
| print(f"Loading checkpoint: {checkpoint_path}") |
| state = torch.load(checkpoint_path, map_location=device, weights_only=False) |
| model.load_state_dict(state) |
| else: |
| print("Evaluating original pre-trained SAM ViT-H (no checkpoint loaded)") |
|
|
| model.eval() |
|
|
| dataset = FacadeDataset(data_dir, split=split, processor=processor, augment=False) |
| dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn) |
|
|
| all_preds = [] |
| all_gts = [] |
| all_ious = [] |
|
|
| print(f"Evaluating on {len(dataset)} samples...") |
| for batch in tqdm(dataloader, desc="Eval"): |
| pixel_values = batch["pixel_values"].to(device) |
| input_boxes = batch["input_boxes"].to(device) |
| gt_masks = batch["ground_truth_mask"].to(device) |
|
|
| outputs = model( |
| pixel_values=pixel_values, |
| input_boxes=input_boxes, |
| multimask_output=False, |
| ) |
|
|
| pred_masks = outputs.pred_masks.squeeze(1).squeeze(1) |
| pred_probs = torch.sigmoid(pred_masks) |
| pred_binary = (pred_probs > 0.5).float() |
|
|
| all_preds.append(pred_binary.cpu().numpy()) |
| all_gts.append(gt_masks.cpu().numpy()) |
|
|
| for b in range(pred_binary.shape[0]): |
| p = pred_binary[b].cpu().numpy() |
| g = gt_masks[b].cpu().numpy() |
| intersection = np.logical_and(p, g).sum() |
| union = np.logical_or(p, g).sum() |
| iou = float(intersection / union) if union > 0 else 1.0 |
| all_ious.append(iou) |
|
|
| all_preds = np.concatenate(all_preds, axis=0) |
| all_gts = np.concatenate(all_gts, axis=0) |
|
|
| metrics = compute_all_metrics(all_preds, all_gts) |
| metrics["mean_iou_list"] = [float(v) for v in all_ious] |
| metrics["split"] = split |
| metrics["checkpoint"] = checkpoint_path if checkpoint_path else "original" |
|
|
| print("\n=== Evaluation Results ===") |
| for k, v in metrics.items(): |
| if isinstance(v, float): |
| print(f"{k}: {v:.4f}") |
|
|
| out_file = os.path.join(output_dir, f"metrics_{split}.json") |
| with open(out_file, "w") as f: |
| json.dump(metrics, f, indent=2) |
|
|
| np.save(os.path.join(output_dir, f"preds_{split}.npy"), all_preds) |
| np.save(os.path.join(output_dir, f"gts_{split}.npy"), all_gts) |
| print(f"Saved to {output_dir}") |
|
|
| return metrics |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--checkpoint", default="", help="Path to fine-tuned checkpoint. Leave empty for original baseline.") |
| parser.add_argument("--data_dir", default="data/cmp_facade") |
| parser.add_argument("--split", default="test") |
| parser.add_argument("--batch_size", type=int, default=2) |
| parser.add_argument("--output_dir", default="outputs/eval") |
| args = parser.parse_args() |
|
|
| evaluate_checkpoint(args.checkpoint, args.data_dir, args.split, args.batch_size, args.output_dir) |
|
|
