File size: 5,954 Bytes
3cc53ab | 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | """
Generate side-by-side qualitative comparisons:
Input image, GT mask, original SAM prediction, fine-tuned SAM prediction.
"""
import os
import argparse
import json
import numpy as np
import torch
from PIL import Image
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from transformers import SamModel, SamProcessor
from tqdm import tqdm
from torch.utils.data import DataLoader
from dataset import FacadeDataset, collate_fn
def get_predictions(model, dataloader, device):
model.eval()
preds = []
with torch.no_grad():
for batch in tqdm(dataloader, desc="Predicting"):
pixel_values = batch["pixel_values"].to(device)
input_boxes = batch["input_boxes"].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_binary = (torch.sigmoid(pred_masks) > 0.5).cpu().numpy()
preds.append(pred_binary)
return np.concatenate(preds, axis=0)
def visualize_comparison(images, gts, preds_baseline, preds_finetuned, indices, save_dir):
os.makedirs(save_dir, exist_ok=True)
for idx in indices:
img = images[idx]
gt = gts[idx]
pred_base = preds_baseline[idx]
pred_ft = preds_finetuned[idx]
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
axes[0, 0].imshow(img)
axes[0, 0].set_title("Input Image")
axes[0, 0].axis('off')
axes[0, 1].imshow(gt, cmap='gray')
axes[0, 1].set_title("Ground Truth")
axes[0, 1].axis('off')
axes[0, 2].axis('off')
iou_base = compute_iou(pred_base, gt)
axes[1, 0].imshow(pred_base, cmap='gray')
axes[1, 0].set_title(f"Original SAM ViT-H\nIoU={iou_base:.3f}")
axes[1, 0].axis('off')
iou_ft = compute_iou(pred_ft, gt)
axes[1, 1].imshow(pred_ft, cmap='gray')
axes[1, 1].set_title(f"Fine-tuned SAM ViT-H\nIoU={iou_ft:.3f}")
axes[1, 1].axis('off')
overlay_base = img.copy()
overlay_base[pred_base > 0] = [255, 0, 0]
blended_base = (img * 0.6 + overlay_base * 0.4).astype(np.uint8)
axes[1, 2].imshow(blended_base)
axes[1, 2].set_title("Overlay Original")
axes[1, 2].axis('off')
plt.suptitle(f"Index {idx}: Original IoU={iou_base:.3f} | Fine-tuned IoU={iou_ft:.3f}", fontsize=14)
plt.tight_layout()
save_path = os.path.join(save_dir, f"comparison_{idx:04d}.png")
plt.savefig(save_path, dpi=150)
plt.close()
print(f"Saved {save_path}")
def compute_iou(pred, gt):
pred = pred.astype(bool)
gt = gt.astype(bool)
inter = np.logical_and(pred, gt).sum()
union = np.logical_or(pred, gt).sum()
return 1.0 if union == 0 else float(inter / union)
def load_raw_images_and_masks(data_dir, split, num_samples):
split_dir = os.path.join(data_dir, split)
with open(os.path.join(split_dir, "metadata.json"), "r") as f:
items = json.load(f)
images = []
gts = []
for item in items[:num_samples]:
img = Image.open(item["image"]).convert("RGB").resize((256, 256), Image.BILINEAR)
mask_path = os.path.join(split_dir, "masks_binary", os.path.basename(item["image"]).replace(".jpg", ".png"))
gt = Image.open(mask_path).convert("L").resize((256, 256), Image.NEAREST)
images.append(np.array(img))
gts.append(np.array(gt) > 0)
return np.array(images), np.array(gts), items
def main(args):
device = "cuda" if torch.cuda.is_available() else "cpu"
processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
dataset = FacadeDataset(args.data_dir, split=args.split, processor=processor, augment=False)
dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, collate_fn=collate_fn)
images_arr, gts_arr, items = load_raw_images_and_masks(args.data_dir, args.split, len(dataset))
print("Running baseline predictions...")
model_base = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
preds_base = get_predictions(model_base, dataloader, device)
del model_base
if torch.cuda.is_available():
torch.cuda.empty_cache()
print("Running fine-tuned predictions...")
model_ft = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
state = torch.load(args.checkpoint, map_location=device, weights_only=False)
model_ft.load_state_dict(state)
preds_ft = get_predictions(model_ft, dataloader, device)
ious_base = []
for i in range(len(preds_base)):
p = preds_base[i]
g = gts_arr[i]
iou = np.logical_and(p, g).sum() / (np.logical_or(p, g).sum() + 1e-6)
ious_base.append(iou)
ious_base = np.array(ious_base)
sorted_idx = np.argsort(ious_base)
selected = [
int(sorted_idx[0]),
int(sorted_idx[-1]),
int(sorted_idx[len(sorted_idx)//2]),
]
rng = np.random.RandomState(42)
extra = rng.choice(len(dataset), size=min(7, max(0, len(dataset)-3)), replace=False)
selected = list(dict.fromkeys(selected + extra.tolist()))[:10]
visualize_comparison(images_arr, gts_arr, preds_base, preds_ft, selected, args.output_dir)
with open(os.path.join(args.output_dir, "comparison_indices.json"), "w") as f:
json.dump({"indices": [int(x) for x in selected], "ious_base": ious_base[selected].tolist()}, f, indent=2)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--checkpoint", required=True)
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/comparison")
args = parser.parse_args()
main(args)
|