| """ |
| PyTorch Dataset for facade segmentation with SAM. |
| Generates bounding-box prompts from ground-truth masks. |
| """ |
| import os |
| import json |
| import random |
| import numpy as np |
| from PIL import Image |
| import torch |
| from torch.utils.data import Dataset |
| from transformers import SamProcessor |
|
|
|
|
| class FacadeDataset(Dataset): |
| """ |
| Dataset for facade segmentation. |
| Loads images and binary masks. Generates bbox prompts from masks. |
| """ |
| def __init__(self, data_dir, split="train", processor=None, augment=False): |
| self.data_dir = data_dir |
| self.split = split |
| self.processor = processor |
| self.augment = augment |
|
|
| split_dir = os.path.join(data_dir, split) |
| with open(os.path.join(split_dir, "metadata.json"), "r") as f: |
| self.items = json.load(f) |
|
|
| self.img_dir = os.path.join(split_dir, "images") |
| self.mask_dir = os.path.join(split_dir, "masks_binary") |
|
|
| def __len__(self): |
| return len(self.items) |
|
|
| def _get_bbox_from_mask(self, mask_np): |
| """Compute bounding box from binary mask with small random jitter.""" |
| ys, xs = np.where(mask_np > 0) |
| if len(xs) == 0: |
| return [0, 0, mask_np.shape[1], mask_np.shape[0]] |
|
|
| x1, y1, x2, y2 = xs.min(), ys.min(), xs.max(), ys.max() |
|
|
| if self.augment: |
| jitter = 10 |
| h, w = mask_np.shape |
| x1 = max(0, x1 + random.randint(-jitter, jitter)) |
| y1 = max(0, y1 + random.randint(-jitter, jitter)) |
| x2 = min(w, x2 + random.randint(-jitter, jitter)) |
| y2 = min(h, y2 + random.randint(-jitter, jitter)) |
|
|
| return [int(x1), int(y1), int(x2), int(y2)] |
|
|
| def __getitem__(self, idx): |
| item = self.items[idx] |
|
|
| img = Image.open(item["image"]).convert("RGB") |
| mask = Image.open(os.path.join(self.mask_dir, os.path.basename(item["mask"]))).convert("L") |
|
|
| mask_np = np.array(mask) |
| bbox = self._get_bbox_from_mask(mask_np) |
|
|
| if self.processor is not None: |
| inputs = self.processor( |
| images=img, |
| input_boxes=[[bbox]], |
| return_tensors="pt", |
| ) |
| for k in inputs: |
| if isinstance(inputs[k], torch.Tensor): |
| inputs[k] = inputs[k].squeeze(0) |
| else: |
| inputs = { |
| "pixel_values": torch.from_numpy(np.array(img)).permute(2, 0, 1).float(), |
| "input_boxes": torch.tensor([bbox], dtype=torch.float32), |
| } |
|
|
| gt_mask = np.array(mask.resize((256, 256), Image.NEAREST)) |
| gt_mask = (gt_mask > 0).astype(np.float32) |
| inputs["ground_truth_mask"] = torch.from_numpy(gt_mask) |
|
|
| return inputs |
|
|
|
|
| def collate_fn(batch): |
| """Collate for DataLoader.""" |
| return { |
| "pixel_values": torch.stack([b["pixel_values"] for b in batch]), |
| "input_boxes": torch.stack([b["input_boxes"] for b in batch]), |
| "ground_truth_mask": torch.stack([b["ground_truth_mask"] for b in batch]), |
| } |
|
|
