ParallelLLC/Segmentation
0
1"""2Data Loader Utilities3 4This module provides data loading utilities for different domains5(satellite, fashion, robotics) with support for few-shot and zero-shot learning.6"""7 8import torch9import torch.nn as nn10import numpy as np11from PIL import Image12import os13import json14from typing import List, Dict, Tuple, Optional15import random16from torch.utils.data import Dataset, DataLoader17import torchvision.transforms as transforms18from torchvision.transforms import functional as F19import cv220 21 22class BaseDataLoader:23 """Base class for domain-specific data loaders."""24 25 def __init__(self, data_dir: str, image_size: Tuple[int, int] = (512, 512)):26 self.data_dir = data_dir27 self.image_size = image_size28 29 # Standard transforms30 self.transform = transforms.Compose([31 transforms.Resize(image_size),32 transforms.ToTensor(),33 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])34 ])35 36 self.mask_transform = transforms.Compose([37 transforms.Resize(image_size, interpolation=transforms.InterpolationMode.NEAREST),38 transforms.ToTensor()39 ])40 41 def load_image(self, image_path: str) -> torch.Tensor:42 """Load and preprocess image."""43 image = Image.open(image_path).convert('RGB')44 return self.transform(image)45 46 def load_mask(self, mask_path: str) -> torch.Tensor:47 """Load and preprocess mask."""48 mask = Image.open(mask_path).convert('L')49 return self.mask_transform(mask)50 51 def get_random_sample(self) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:52 """Get a random sample from the dataset."""53 raise NotImplementedError54 55 def get_class_examples(self, class_name: str, num_examples: int) -> List[Tuple[torch.Tensor, torch.Tensor]]:56 """Get examples for a specific class."""57 raise NotImplementedError58 59 60class SatelliteDataLoader(BaseDataLoader):61 """Data loader for satellite imagery segmentation."""62 63 def __init__(self, data_dir: str, image_size: Tuple[int, int] = (512, 512)):64 super().__init__(data_dir, image_size)65 66 # Satellite-specific classes67 self.classes = ["building", "road", "vegetation", "water"]68 self.class_to_id = {cls: i for i, cls in enumerate(self.classes)}69 70 # Load dataset structure71 self.load_dataset_structure()72 73 def load_dataset_structure(self):74 """Load dataset structure and file paths."""75 self.images = []76 self.masks = []77 self.class_samples = {cls: [] for cls in self.classes}78 79 # Assuming structure: data_dir/images/ and data_dir/masks/80 images_dir = os.path.join(self.data_dir, "images")81 masks_dir = os.path.join(self.data_dir, "masks")82 83 if not os.path.exists(images_dir) or not os.path.exists(masks_dir):84 # Create dummy data for demonstration85 self.create_dummy_data()86 return87 88 # Load real data89 for filename in os.listdir(images_dir):90 if filename.endswith(('.jpg', '.png', '.tif')):91 image_path = os.path.join(images_dir, filename)92 mask_path = os.path.join(masks_dir, filename.replace('.jpg', '_mask.png'))93 94 if os.path.exists(mask_path):95 self.images.append(image_path)96 self.masks.append(mask_path)97 98 # Categorize by class (simplified)99 self.categorize_sample(image_path, mask_path)100 101 def create_dummy_data(self):102 """Create dummy satellite data for demonstration."""103 print("Creating dummy satellite data...")104 105 # Create dummy directory structure106 os.makedirs(os.path.join(self.data_dir, "images"), exist_ok=True)107 os.makedirs(os.path.join(self.data_dir, "masks"), exist_ok=True)108 109 # Generate dummy images and masks110 for i in range(100):111 # Create dummy image (satellite-like)112 image = np.random.randint(50, 200, (512, 512, 3), dtype=np.uint8)113 114 # Add some structure to make it look like satellite imagery115 # Buildings (rectangular shapes)116 for _ in range(5):117 x, y = np.random.randint(0, 400), np.random.randint(0, 400)118 w, h = np.random.randint(20, 80), np.random.randint(20, 80)119 image[y:y+h, x:x+w] = np.random.randint(100, 150, 3)120 121 # Roads (linear structures)122 for _ in range(3):123 x, y = np.random.randint(0, 512), np.random.randint(0, 512)124 length = np.random.randint(50, 150)125 angle = np.random.uniform(0, 2*np.pi)126 for j in range(length):127 px = int(x + j * np.cos(angle))128 py = int(y + j * np.sin(angle))129 if 0 <= px < 512 and 0 <= py < 512:130 image[py, px] = [80, 80, 80]131 132 # Save image133 image_path = os.path.join(self.data_dir, "images", f"satellite_{i:03d}.jpg")134 Image.fromarray(image).save(image_path)135 136 # Create corresponding mask137 mask = np.zeros((512, 512), dtype=np.uint8)138 139 # Add building masks140 for _ in range(3):141 x, y = np.random.randint(0, 400), np.random.randint(0, 400)142 w, h = np.random.randint(20, 80), np.random.randint(20, 80)143 mask[y:y+h, x:x+w] = 1 # Building class144 145 # Add road masks146 for _ in range(2):147 x, y = np.random.randint(0, 512), np.random.randint(0, 512)148 length = np.random.randint(50, 150)149 angle = np.random.uniform(0, 2*np.pi)150 for j in range(length):151 px = int(x + j * np.cos(angle))152 py = int(y + j * np.sin(angle))153 if 0 <= px < 512 and 0 <= py < 512:154 mask[py, px] = 2 # Road class155 156 # Save mask157 mask_path = os.path.join(self.data_dir, "masks", f"satellite_{i:03d}_mask.png")158 Image.fromarray(mask * 85).save(mask_path) # Scale for visibility159 160 # Add to lists161 self.images.append(image_path)162 self.masks.append(mask_path)163 164 # Categorize165 self.categorize_sample(image_path, mask_path)166 167 def categorize_sample(self, image_path: str, mask_path: str):168 """Categorize sample by dominant class."""169 mask = np.array(Image.open(mask_path))170 171 # Count pixels for each class172 class_counts = {}173 for i, class_name in enumerate(self.classes):174 class_counts[class_name] = np.sum(mask == i)175 176 # Find dominant class177 dominant_class = max(class_counts.items(), key=lambda x: x[1])[0]178 self.class_samples[dominant_class].append((image_path, mask_path))179 180 def get_random_query(self, class_name: str) -> Tuple[torch.Tensor, torch.Tensor]:181 """Get a random query image and mask for a specific class."""182 if class_name not in self.class_samples or not self.class_samples[class_name]:183 # Fallback to any available sample184 idx = random.randint(0, len(self.images) - 1)185 image = self.load_image(self.images[idx])186 mask = self.load_mask(self.masks[idx])187 return image, mask188 189 # Get random sample from specified class190 image_path, mask_path = random.choice(self.class_samples[class_name])191 image = self.load_image(image_path)192 mask = self.load_mask(mask_path)193 194 return image, mask195 196 def get_class_examples(self, class_name: str, num_examples: int) -> List[Tuple[torch.Tensor, torch.Tensor]]:197 """Get examples for a specific class."""198 examples = []199 200 if class_name in self.class_samples:201 available_samples = self.class_samples[class_name]202 selected_samples = random.sample(available_samples, min(num_examples, len(available_samples)))203 204 for image_path, mask_path in selected_samples:205 image = self.load_image(image_path)206 mask = self.load_mask(mask_path)207 examples.append((image, mask))208 209 return examples210 211 212class FashionDataLoader(BaseDataLoader):213 """Data loader for fashion segmentation."""214 215 def __init__(self, data_dir: str, image_size: Tuple[int, int] = (512, 512)):216 super().__init__(data_dir, image_size)217 218 # Fashion-specific classes219 self.classes = ["shirt", "pants", "dress", "shoes"]220 self.class_to_id = {cls: i for i, cls in enumerate(self.classes)}221 222 # Load dataset structure223 self.load_dataset_structure()224 225 def load_dataset_structure(self):226 """Load dataset structure and file paths."""227 self.images = []228 self.masks = []229 self.class_samples = {cls: [] for cls in self.classes}230 231 # Assuming structure: data_dir/images/ and data_dir/masks/232 images_dir = os.path.join(self.data_dir, "images")233 masks_dir = os.path.join(self.data_dir, "masks")234 235 if not os.path.exists(images_dir) or not os.path.exists(masks_dir):236 # Create dummy data for demonstration237 self.create_dummy_data()238 return239 240 # Load real data241 for filename in os.listdir(images_dir):242 if filename.endswith(('.jpg', '.png')):243 image_path = os.path.join(images_dir, filename)244 mask_path = os.path.join(masks_dir, filename.replace('.jpg', '_mask.png'))245 246 if os.path.exists(mask_path):247 self.images.append(image_path)248 self.masks.append(mask_path)249 250 # Categorize by class251 self.categorize_sample(image_path, mask_path)252 253 def create_dummy_data(self):254 """Create dummy fashion data for demonstration."""255 print("Creating dummy fashion data...")256 257 # Create dummy directory structure258 os.makedirs(os.path.join(self.data_dir, "images"), exist_ok=True)259 os.makedirs(os.path.join(self.data_dir, "masks"), exist_ok=True)260 261 # Generate dummy images and masks262 for i in range(100):263 # Create dummy image (fashion-like)264 image = np.random.randint(200, 255, (512, 512, 3), dtype=np.uint8)265 266 # Add fashion items267 class_id = i % len(self.classes)268 269 if class_id == 0: # Shirt270 # Create shirt-like shape271 center_x, center_y = 256, 256272 width, height = 150, 200273 image[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = [100, 150, 200]274 275 elif class_id == 1: # Pants276 # Create pants-like shape277 center_x, center_y = 256, 300278 width, height = 120, 180279 image[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = [50, 100, 150]280 281 elif class_id == 2: # Dress282 # Create dress-like shape283 center_x, center_y = 256, 250284 width, height = 140, 220285 image[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = [200, 100, 150]286 287 else: # Shoes288 # Create shoes-like shape289 center_x, center_y = 256, 400290 width, height = 100, 60291 image[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = [80, 80, 80]292 293 # Save image294 image_path = os.path.join(self.data_dir, "images", f"fashion_{i:03d}.jpg")295 Image.fromarray(image).save(image_path)296 297 # Create corresponding mask298 mask = np.zeros((512, 512), dtype=np.uint8)299 300 # Add mask for the fashion item301 if class_id == 0: # Shirt302 center_x, center_y = 256, 256303 width, height = 150, 200304 mask[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = 1305 306 elif class_id == 1: # Pants307 center_x, center_y = 256, 300308 width, height = 120, 180309 mask[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = 2310 311 elif class_id == 2: # Dress312 center_x, center_y = 256, 250313 width, height = 140, 220314 mask[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = 3315 316 else: # Shoes317 center_x, center_y = 256, 400318 width, height = 100, 60319 mask[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = 4320 321 # Save mask322 mask_path = os.path.join(self.data_dir, "masks", f"fashion_{i:03d}_mask.png")323 Image.fromarray(mask * 51).save(mask_path) # Scale for visibility324 325 # Add to lists326 self.images.append(image_path)327 self.masks.append(mask_path)328 329 # Categorize330 self.categorize_sample(image_path, mask_path)331 332 def categorize_sample(self, image_path: str, mask_path: str):333 """Categorize sample by dominant class."""334 mask = np.array(Image.open(mask_path))335 336 # Count pixels for each class337 class_counts = {}338 for i, class_name in enumerate(self.classes):339 class_counts[class_name] = np.sum(mask == (i + 1)) # +1 because 0 is background340 341 # Find dominant class342 dominant_class = max(class_counts.items(), key=lambda x: x[1])[0]343 self.class_samples[dominant_class].append((image_path, mask_path))344 345 def get_test_sample(self) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:346 """Get a random test sample with ground truth masks."""347 idx = random.randint(0, len(self.images) - 1)348 image = self.load_image(self.images[idx])349 mask = self.load_mask(self.masks[idx])350 351 # Convert single mask to multi-class dictionary352 ground_truth = {}353 for i, class_name in enumerate(self.classes):354 class_mask = (mask == (i + 1)).float() # +1 because 0 is background355 ground_truth[class_name] = class_mask356 357 return image, ground_truth358 359 360class RoboticsDataLoader(BaseDataLoader):361 """Data loader for robotics segmentation."""362 363 def __init__(self, data_dir: str, image_size: Tuple[int, int] = (512, 512)):364 super().__init__(data_dir, image_size)365 366 # Robotics-specific classes367 self.classes = ["robot", "tool", "safety"]368 self.class_to_id = {cls: i for i, cls in enumerate(self.classes)}369 370 # Load dataset structure371 self.load_dataset_structure()372 373 def load_dataset_structure(self):374 """Load dataset structure and file paths."""375 self.images = []376 self.masks = []377 self.class_samples = {cls: [] for cls in self.classes}378 379 # Assuming structure: data_dir/images/ and data_dir/masks/380 images_dir = os.path.join(self.data_dir, "images")381 masks_dir = os.path.join(self.data_dir, "masks")382 383 if not os.path.exists(images_dir) or not os.path.exists(masks_dir):384 # Create dummy data for demonstration385 self.create_dummy_data()386 return387 388 # Load real data389 for filename in os.listdir(images_dir):390 if filename.endswith(('.jpg', '.png')):391 image_path = os.path.join(images_dir, filename)392 mask_path = os.path.join(masks_dir, filename.replace('.jpg', '_mask.png'))393 394 if os.path.exists(mask_path):395 self.images.append(image_path)396 self.masks.append(mask_path)397 398 # Categorize by class399 self.categorize_sample(image_path, mask_path)400 401 def create_dummy_data(self):402 """Create dummy robotics data for demonstration."""403 print("Creating dummy robotics data...")404 405 # Create dummy directory structure406 os.makedirs(os.path.join(self.data_dir, "images"), exist_ok=True)407 os.makedirs(os.path.join(self.data_dir, "masks"), exist_ok=True)408 409 # Generate dummy images and masks410 for i in range(100):411 # Create dummy image (robotics-like)412 image = np.random.randint(50, 150, (512, 512, 3), dtype=np.uint8)413 414 # Add robotics elements415 class_id = i % len(self.classes)416 417 if class_id == 0: # Robot418 # Create robot-like shape419 center_x, center_y = 256, 256420 width, height = 120, 160421 image[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = [100, 100, 100]422 423 elif class_id == 1: # Tool424 # Create tool-like shape425 center_x, center_y = 256, 256426 width, height = 80, 120427 image[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = [150, 100, 50]428 429 else: # Safety equipment430 # Create safety equipment-like shape431 center_x, center_y = 256, 256432 width, height = 100, 100433 image[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = [200, 200, 50]434 435 # Save image436 image_path = os.path.join(self.data_dir, "images", f"robotics_{i:03d}.jpg")437 Image.fromarray(image).save(image_path)438 439 # Create corresponding mask440 mask = np.zeros((512, 512), dtype=np.uint8)441 442 # Add mask for the robotics element443 if class_id == 0: # Robot444 center_x, center_y = 256, 256445 width, height = 120, 160446 mask[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = 1447 448 elif class_id == 1: # Tool449 center_x, center_y = 256, 256450 width, height = 80, 120451 mask[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = 2452 453 else: # Safety equipment454 center_x, center_y = 256, 256455 width, height = 100, 100456 mask[center_y-height//2:center_y+height//2, center_x-width//2:center_x+width//2] = 3457 458 # Save mask459 mask_path = os.path.join(self.data_dir, "masks", f"robotics_{i:03d}_mask.png")460 Image.fromarray(mask * 85).save(mask_path) # Scale for visibility461 462 # Add to lists463 self.images.append(image_path)464 self.masks.append(mask_path)465 466 # Categorize467 self.categorize_sample(image_path, mask_path)468 469 def categorize_sample(self, image_path: str, mask_path: str):470 """Categorize sample by dominant class."""471 mask = np.array(Image.open(mask_path))472 473 # Count pixels for each class474 class_counts = {}475 for i, class_name in enumerate(self.classes):476 class_counts[class_name] = np.sum(mask == (i + 1)) # +1 because 0 is background477 478 # Find dominant class479 dominant_class = max(class_counts.items(), key=lambda x: x[1])[0]480 self.class_samples[dominant_class].append((image_path, mask_path))481 482 def get_test_sample(self) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:483 """Get a random test sample with ground truth masks."""484 idx = random.randint(0, len(self.images) - 1)485 image = self.load_image(self.images[idx])486 mask = self.load_mask(self.masks[idx])487 488 # Convert single mask to multi-class dictionary489 ground_truth = {}490 for i, class_name in enumerate(self.classes):491 class_mask = (mask == (i + 1)).float() # +1 because 0 is background492 ground_truth[class_name] = class_mask493 494 return image, ground_truth