| """ |
| Inference pipeline for fine-tuned SAM ViT-H facade segmentation. |
| """ |
| import os |
| import argparse |
| 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 |
|
|
|
|
| def load_model(checkpoint_path=None, device=None): |
| if device is None: |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
| 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): |
| state = torch.load(checkpoint_path, map_location=device, weights_only=False) |
| model.load_state_dict(state) |
| print(f"Loaded checkpoint from {checkpoint_path}") |
| else: |
| print("Using original pre-trained SAM ViT-H") |
|
|
| model.eval() |
| return model, processor, device |
|
|
|
|
| def predict_facade(model, processor, image_path, bbox=None, device="cpu"): |
| img = Image.open(image_path).convert("RGB") |
| orig_w, orig_h = img.size |
|
|
| if bbox is None: |
| bbox = [[[0.0, 0.0, float(orig_w), float(orig_h)]]] |
| else: |
| bbox = [[[float(b) for b in bbox]]] |
|
|
| inputs = processor(images=img, input_boxes=bbox, return_tensors="pt").to(device) |
|
|
| with torch.no_grad(): |
| outputs = model( |
| pixel_values=inputs["pixel_values"], |
| input_boxes=inputs["input_boxes"], |
| multimask_output=False, |
| ) |
|
|
| pred_mask = outputs.pred_masks.squeeze().cpu().numpy() |
| pred_binary = (pred_mask > 0).astype(np.uint8) * 255 |
| pred_binary = Image.fromarray(pred_binary).resize((orig_w, orig_h), Image.NEAREST) |
| return np.array(pred_binary), img |
|
|
|
|
| def visualize_result(image, mask, save_path): |
| fig, axes = plt.subplots(1, 3, figsize=(15, 5)) |
| img_arr = np.array(image) |
|
|
| axes[0].imshow(img_arr) |
| axes[0].set_title("Input Image") |
| axes[0].axis('off') |
|
|
| axes[1].imshow(mask, cmap='gray') |
| axes[1].set_title("Predicted Facade Mask") |
| axes[1].axis('off') |
|
|
| overlay = img_arr.copy() |
| overlay[mask > 0] = [255, 0, 0] |
| blended = (img_arr * 0.6 + overlay * 0.4).astype(np.uint8) |
| axes[2].imshow(blended) |
| axes[2].set_title("Overlay") |
| axes[2].axis('off') |
|
|
| plt.tight_layout() |
| plt.savefig(save_path, dpi=150) |
| plt.close() |
| print(f"Result saved to {save_path}") |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--checkpoint", default="", help="Path to fine-tuned checkpoint") |
| parser.add_argument("--image", required=True, help="Path to input image") |
| parser.add_argument("--bbox", default=None, type=str, help="Bounding box as x1,y1,x2,y2") |
| parser.add_argument("--output", default="outputs/inference_result.png", help="Output path") |
| args = parser.parse_args() |
|
|
| model, processor, device = load_model(args.checkpoint) |
|
|
| bbox = None |
| if args.bbox: |
| bbox = [float(x.strip()) for x in args.bbox.split(",")] |
| assert len(bbox) == 4, "Bounding box must be x1,y1,x2,y2" |
|
|
| mask, img = predict_facade(model, processor, args.image, bbox, device) |
|
|
| os.makedirs(os.path.dirname(args.output), exist_ok=True) |
| visualize_result(img, mask, args.output) |
|
|
| mask_path = args.output.replace('.png', '_mask.png') |
| Image.fromarray(mask).save(mask_path) |
| print(f"Mask saved to {mask_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
