Source code for aitlas.datasets.landcover_ai

import glob
import os
import cv2
import numpy as np

from ..utils import image_loader
from .semantic_segmentation import SemanticSegmentationDataset

"""
41 orthophoto tiles from different counties located in all regions. Every tile has about 5 km2.
There are 33 images with resolution 25cm (ca. 9000 × 9500 px) and 8 images with resolution 50cm (ca. 4200 × 4700 px)
Tne masks are codded with building (1), woodland (2), water (3), and road (4)
Use function split_images to split the images and the masks in smaller patches
"""


[docs]class LandCoverAiDataset(SemanticSegmentationDataset): url = "https://landcover.ai.linuxpolska.com/" labels = ["Background", "Buildings", "Woodlands", "Water", "Road"] color_mapping = [[255, 255, 0], [0, 0, 0], [0, 255, 0], [0, 0, 255], [200, 200, 200]] name = "Landcover AI" def __init__(self, config): # now call the constructor to validate the schema and split the data super().__init__(config) def __getitem__(self, index): image = image_loader(self.images[index]) mask = image_loader(self.masks[index])[:, :, 1] # extract certain classes from mask (e.g. Buildings) masks = [(mask == v) for v, label in enumerate(self.labels)] mask = np.stack(masks, axis=-1).astype("float32") return self.apply_transformations(image, mask)
[docs]def split_images(imgs_dir, masks_dir, output_dir): target_size = 512 img_paths = glob.glob(os.path.join(imgs_dir, "*.tif")) mask_paths = glob.glob(os.path.join(masks_dir, "*.tif")) img_paths.sort() mask_paths.sort() os.makedirs(output_dir) for i, (img_path, mask_path) in enumerate(zip(img_paths, mask_paths)): img_filename = os.path.splitext(os.path.basename(img_path))[0] mask_filename = os.path.splitext(os.path.basename(mask_path))[0] img = cv2.imread(img_path) mask = cv2.imread(mask_path) assert img_filename == mask_filename and img.shape[:2] == mask.shape[:2] k = 0 for y in range(0, img.shape[0], target_size): for x in range(0, img.shape[1], target_size): img_tile = img[y: y + target_size, x: x + target_size] mask_tile = mask[y: y + target_size, x: x + target_size] if ( img_tile.shape[0] == target_size and img_tile.shape[1] == target_size ): out_img_path = os.path.join( output_dir, "{}_{}.jpg".format(img_filename, k) ) cv2.imwrite(out_img_path, img_tile) out_mask_path = os.path.join( output_dir, "{}_{}_m.png".format(mask_filename, k) ) cv2.imwrite(out_mask_path, mask_tile) k += 1 print("Processed {} {}/{}".format(img_filename, i + 1, len(img_paths)))