""" Evaluation Script for Text-Conditioned Segmentation Model Computes mIoU, Dice Score, and generates visual examples """ import torch import numpy as np from PIL import Image import json import os from pathlib import Path from tqdm import tqdm from model import SegModel from inference import predict, PROMPT_TO_MODE from dataset import SegmentationDataset from pycocotools.coco import COCO DEVICE = "cpu" MODEL_PATH = "best_model.pth" RESULTS_DIR = "../reports" os.makedirs(RESULTS_DIR, exist_ok=True) os.makedirs(f"{RESULTS_DIR}/visuals", exist_ok=True) def compute_iou(pred, gt, threshold=0.5): """Compute Intersection over Union""" pred_binary = (pred > threshold).astype(np.float32) gt_binary = (gt > 0).astype(np.float32) intersection = np.sum(pred_binary * gt_binary) union = np.sum(pred_binary) + np.sum(gt_binary) - intersection if union == 0: return 1.0 if np.sum(pred_binary) == 0 and np.sum(gt_binary) == 0 else 0.0 return intersection / union def compute_dice(pred, gt, threshold=0.5): """Compute Dice Score (F1 for binary segmentation)""" pred_binary = (pred > threshold).astype(np.float32) gt_binary = (gt > 0).astype(np.float32) intersection = 2.0 * np.sum(pred_binary * gt_binary) total = np.sum(pred_binary) + np.sum(gt_binary) if total == 0: return 1.0 if np.sum(pred_binary) == 0 and np.sum(gt_binary) == 0 else 0.0 return intersection / total def load_ground_truth(image_id, dataset_path, coco_ann_file): """Load ground truth mask from COCO annotation""" try: coco = COCO(coco_ann_file) img_info = coco.loadImgs(image_id) if not img_info: return None ann_ids = coco.getAnnIds(imgIds=image_id) anns = coco.loadAnns(ann_ids) if not anns: return None # Create binary mask from all annotations img_data = img_info[0] mask = np.zeros((img_data['height'], img_data['width']), dtype=np.uint8) for ann in anns: rle = ann['segmentation'] m = coco.annToMask(ann) mask = np.maximum(mask, m) return mask * 255 except Exception as e: print(f"Error loading GT for image {image_id}: {e}") return None def evaluate_dataset(dataset_type="crack", split="valid"): """Evaluate on a specific dataset and split""" print(f"\n{'='*70}") print(f"Evaluating {dataset_type.upper()} - {split.upper()} Set") print(f"{'='*70}") # Determine dataset paths if dataset_type == "crack": dataset_root = "../data/cracks.v1-cracks-f.coco" prompts = ["segment crack", "segment wall crack"] else: dataset_root = "../data/Drywall-Join-Detect.v2i.coco" prompts = ["segment taping area", "segment joint", "segment drywall seam"] coco_ann_file = f"{dataset_root}/{split}/_annotations.coco.json" image_dir = f"{dataset_root}/{split}" if not os.path.exists(coco_ann_file): print(f"⚠️ Annotation file not found: {coco_ann_file}") return None # Load COCO dataset coco = COCO(coco_ann_file) img_ids = coco.getImgIds() results = { "dataset": dataset_type, "split": split, "total_images": len(img_ids), "prompts": {}, "timestamp": None } # Test each prompt for prompt in prompts: print(f"\n Testing prompt: {prompt}") ious = [] dices = [] visual_samples = [] for idx, img_id in enumerate(tqdm(img_ids[:50], desc=prompt)): # Limit to first 50 for speed try: img_info = coco.loadImgs(img_id)[0] image_path = os.path.join(image_dir, img_info['file_name']) if not os.path.exists(image_path): continue # Run inference img_orig, pred_mask = predict(image_path, prompt, MODEL_PATH) pred_array = np.array(pred_mask) # Load ground truth gt_mask = load_ground_truth(img_id, image_dir, coco_ann_file) if gt_mask is None: continue # Resize GT to match prediction gt_mask = Image.fromarray(gt_mask).resize( (pred_array.shape[1], pred_array.shape[0]), Image.Resampling.BILINEAR ) gt_array = np.array(gt_mask) # Compute metrics iou = compute_iou(pred_array, gt_array) dice = compute_dice(pred_array, gt_array) ious.append(iou) dices.append(dice) # Store visual samples (first 2 successful predictions) if len(visual_samples) < 2: visual_samples.append({ "image_id": img_id, "image_path": image_path, "original": img_orig, "gt": Image.fromarray(gt_array), "pred": Image.fromarray(pred_array), "iou": iou, "dice": dice }) except Exception as e: continue # Compute statistics if ious: miou = np.mean(ious) dice_mean = np.mean(dices) results["prompts"][prompt] = { "mIoU": float(miou), "Dice": float(dice_mean), "samples_evaluated": len(ious), "visual_samples": visual_samples } print(f" ✅ mIoU: {miou:.4f}") print(f" ✅ Dice: {dice_mean:.4f}") print(f" ✅ Samples: {len(ious)}") return results def create_visual_comparison(results): """Create side-by-side comparison images""" print(f"\n{'='*70}") print("Creating Visual Examples") print(f"{'='*70}") for dataset_type in ["crack", "crack"]: # Only evaluate what we have dataset_results = evaluate_dataset(dataset_type) if not dataset_results: continue for prompt, metrics in dataset_results["prompts"].items(): for idx, sample in enumerate(metrics.get("visual_samples", [])): # Create comparison image orig = sample["original"] gt = sample["gt"] pred = sample["pred"] # Resize all to same size w, h = orig.size gt = gt.resize((w, h), Image.Resampling.NEAREST) pred = pred.resize((w, h), Image.Resampling.NEAREST) # Create 3-panel comparison comparison = Image.new('RGB', (w*3 + 20, h + 40), color='white') # Convert grayscale to RGB for display orig_rgb = orig.convert('RGB') gt_rgb = Image.new('RGB', gt.size, color='white') gt_rgb.paste(gt) pred_rgb = Image.new('RGB', pred.size, color='white') pred_rgb.paste(pred) comparison.paste(orig_rgb, (0, 40)) comparison.paste(gt_rgb, (w + 10, 40)) comparison.paste(pred_rgb, (w*2 + 20, 40)) # Add labels from PIL import ImageDraw, ImageFont draw = ImageDraw.Draw(comparison) labels = [ f"Original (ID: {sample['image_id']})", f"Ground Truth", f"Prediction (IoU: {sample['iou']:.3f})" ] for i, label in enumerate(labels): x = i * (w + 10) + 5 draw.text((x, 5), label, fill='black') # Save comparison prompt_slug = prompt.replace(" ", "_").lower() save_path = f"{RESULTS_DIR}/visuals/{dataset_type}__{prompt_slug}__{idx}.png" comparison.save(save_path) print(f" Saved: {save_path}") def generate_report(): """Generate complete evaluation report""" print("\n" + "="*70) print("STARTING COMPREHENSIVE EVALUATION") print("="*70) # Evaluate both datasets crack_results = evaluate_dataset("crack", "valid") drywall_results = evaluate_dataset("drywall", "valid") # Generate report report = { "title": "Text-Conditioned Segmentation Evaluation Report", "model": "ResNet18 Encoder + UNet Decoder", "device": str(DEVICE), "datasets": { "crack": crack_results, "drywall": drywall_results } } # Save JSON report report_json = f"{RESULTS_DIR}/evaluation_metrics.json" with open(report_json, 'w') as f: # Convert non-serializable objects clean_report = { "title": report["title"], "model": report["model"], "device": report["device"], "datasets": {} } for ds_type, ds_data in report["datasets"].items(): if ds_data: clean_report["datasets"][ds_type] = { "dataset": ds_data["dataset"], "split": ds_data["split"], "total_images": ds_data["total_images"], "prompts": { prompt: { "mIoU": metrics["mIoU"], "Dice": metrics["Dice"], "samples_evaluated": metrics["samples_evaluated"] } for prompt, metrics in ds_data["prompts"].items() } } json.dump(clean_report, f, indent=2) print(f"\n✅ Report saved to: {report_json}") # Create visuals create_visual_comparison(report) return report if __name__ == "__main__": report = generate_report() print("\n" + "="*70) print("EVALUATION COMPLETE") print("="*70) print(f"\nResults saved to: {RESULTS_DIR}/")