Spaces:
Build error
Build error
| import os | |
| from PIL import Image | |
| import numpy as np | |
| # Define paths for the dataset folders | |
| IMAGES_DIR = "miniFrance/labeled_training/labeled" | |
| LABELS_DIR = "miniFrance/labels" | |
| # Output directories for patches | |
| OUTPUT_IMAGES_DIR = "output2/patches/images" | |
| OUTPUT_LABELS_DIR = "output2/patches/labels" | |
| # Create output directories if they do not exist | |
| os.makedirs(OUTPUT_IMAGES_DIR, exist_ok=True) | |
| os.makedirs(OUTPUT_LABELS_DIR, exist_ok=True) | |
| # Define label color mapping (0-15) | |
| label_color_mapping = { | |
| 0: (0, 0, 0), # No information | |
| 1: (128, 64, 128), # Urban fabric | |
| 2: (244, 35, 232), # Industrial, commercial, public, military, private and transport units | |
| 3: (70, 70, 70), # Mine, dump and construction sites | |
| 4: (102, 102, 156), # Artificial non-agricultural vegetated areas | |
| 5: (190, 153, 153), # Arable land (annual crops) | |
| 6: (153, 153, 153), # Permanent crops | |
| 7: (250, 170, 30), # Pastures | |
| 8: (220, 220, 0), # Complex and mixed cultivation patterns | |
| 9: (107, 142, 35), # Orchards at the fringe of urban classes | |
| 10: (152, 251, 152), # Forests | |
| 11: (70, 130, 180), # Herbaceous vegetation associations | |
| 12: (220, 20, 60), # Open spaces with little or no vegetation | |
| 13: (255, 0, 0), # Wetlands | |
| 14: (0, 0, 142), # Water | |
| 15: (220, 220, 220) # Clouds and shadows | |
| } | |
| def convert_label_to_rgb(label_array, mapping): | |
| """ | |
| Convert a 2D array of label indices into a 3D RGB image using the provided mapping. | |
| """ | |
| h, w = label_array.shape | |
| rgb_image = np.zeros((h, w, 3), dtype=np.uint8) | |
| for label, color in mapping.items(): | |
| mask = (label_array == label) | |
| rgb_image[mask] = color | |
| return rgb_image | |
| def get_base_key(filename): | |
| """ | |
| Extract the full base key from the filename. | |
| For example: | |
| - "06-2014-1015-6335-LA93-0M50-E080.jp2.tif" becomes "06-2014-1015-6335-LA93-0M50-E080" | |
| - "06-2014-1015-6335-LA93-0M50-E080_UA2012.tif" becomes "06-2014-1015-6335-LA93-0M50-E080" | |
| """ | |
| base = os.path.basename(filename) | |
| base_no_ext = base.split('.')[0] # Remove compound extension | |
| if "_" in base_no_ext: | |
| key = base_no_ext.split("_")[0] | |
| else: | |
| key = base_no_ext | |
| return key | |
| # Build a dictionary for label files by base key for fast lookup. | |
| label_files_dict = {} | |
| for root, dirs, files in os.walk(LABELS_DIR): | |
| for file in files: | |
| if file.lower().endswith((".tif", ".tiff")): | |
| key = get_base_key(file) | |
| full_path = os.path.join(root, file) | |
| label_files_dict.setdefault(key, []).append(full_path) | |
| # Process each image file in the images directory | |
| patch_size = 1000 # size of the patch (1000 x 1000) | |
| for root, dirs, files in os.walk(IMAGES_DIR): | |
| for file in files: | |
| if file.lower().endswith((".tif", ".tiff", ".jp2.tif")): | |
| image_path = os.path.join(root, file) | |
| key = get_base_key(file) | |
| # Check if a corresponding label exists by base key | |
| if key not in label_files_dict: | |
| continue # skip if no matching label file | |
| # For simplicity, take the first matching label file (if there are multiples) | |
| label_path = label_files_dict[key][0] | |
| # Open the image and label | |
| try: | |
| image = Image.open(image_path) | |
| label = Image.open(label_path) | |
| except Exception as e: | |
| print(f"Error opening {image_path} or {label_path}: {e}") | |
| continue | |
| # Ensure both image and label have same dimensions | |
| if image.size != label.size: | |
| print(f"Size mismatch between image {image_path} and label {label_path}. Skipping.") | |
| continue | |
| width, height = image.size | |
| # Split into non-overlapping patches of size patch_size x patch_size | |
| num_patches_x = width // patch_size | |
| num_patches_y = height // patch_size | |
| for i in range(num_patches_x): | |
| for j in range(num_patches_y): | |
| left = i * patch_size | |
| upper = j * patch_size | |
| right = left + patch_size | |
| lower = upper + patch_size | |
| # Crop the image and label patch | |
| image_patch = image.crop((left, upper, right, lower)) | |
| label_patch = label.crop((left, upper, right, lower)) | |
| # Convert label patch to a numpy array and then to RGB | |
| label_array = np.array(label_patch) | |
| # Check if the patch has at least 2 unique labels | |
| unique_labels = np.unique(label_array) | |
| if len(unique_labels) < 2: | |
| continue # Skip patches with less than 2 unique labels | |
| rgb_label_patch = convert_label_to_rgb(label_array, label_color_mapping) | |
| rgb_label_patch_img = Image.fromarray(rgb_label_patch) | |
| # Generate patch filename using original name and patch indices | |
| base_image_name = os.path.splitext(os.path.basename(file))[0] | |
| patch_filename = f"{base_image_name}_patch_{i}_{j}.png" | |
| # Resize images to 256x256 before saving | |
| image_patch_resized = image_patch.resize((256, 256), Image.BILINEAR) | |
| rgb_label_patch_img_resized = rgb_label_patch_img.resize((256, 256), Image.NEAREST) | |
| # Save the resized image and label patches | |
| image_patch_resized.save(os.path.join(OUTPUT_IMAGES_DIR, patch_filename)) | |
| rgb_label_patch_img_resized.save(os.path.join(OUTPUT_LABELS_DIR, patch_filename)) | |
| print(f"Processed {file} with corresponding label {os.path.basename(label_path)}") | |