image_output_5.png image_output_7.png image_output_4.png image_output_6.png image_output_3.png image_output_2.png image_output_1.png
from PIL import Image

from torch.utils.data import random_split, DataLoader, Dataset

from torchvision.io import read_image, ImageReadMode

import matplotlib.pyplot as plt

import os

import torchvision.transforms.functional as TF

BATCH_SIZE = 16
IMAGE_SIZE = (256, 256)
IN_CHANNELS = 3
LEARNING_RATE = 0.0001
IMAGES_DIR = '/kaggle/input/danish-golf-courses-orthophotos/1. orthophotos/'
SEGMASKS_DIR = '/kaggle/input/danish-golf-courses-orthophotos/2. segmentation masks/'
LABELMASKS_DIR = '/kaggle/input/danish-golf-courses-orthophotos/3. class masks/'

#Loading the data
orthophoto_list = os.listdir(IMAGES_DIR)
print("There are ", len(orthophoto_list), " orthophotos in this dataset!")

#Load image with index of 5 (I prefer this image as it shows all the classes)
idx = 5 #The index can be changed to view other orthophotos.
golf_image = Image.open(os.path.join(IMAGES_DIR, orthophoto_list[idx]))
golf_segmask = Image.open(os.path.join(SEGMASKS_DIR, orthophoto_list[idx].replace(".jpg", ".png"))) #The class masks are png instead of jpg

#Plot using matplotlib
fig, axes = plt.subplots(1, 2)

axes[0].set_title('Orthophoto')
axes[1].set_title('Segmentation Mask')

axes[0].imshow(golf_image)
axes[1].imshow(golf_segmask)

class GolfDataset(Dataset):
    def __init__(self, images_dir, labelmasks_dir):
        self.images_dir = images_dir
        self.labelmasks_dir = labelmasks_dir
        self.images_dir_list = os.listdir(images_dir)
    def __len__(self):
        return len(self.images_dir_list)
    def __getitem__(self, idx):
        image_path = os.path.join(self.images_dir, self.images_dir_list[idx])
        image = read_image(image_path, ImageReadMode.RGB)
        label_mask_path = os.path.join(self.labelmasks_dir, self.images_dir_list[idx]).replace('.jpg', '.png')
        label_mask = read_image(label_mask_path, ImageReadMode.GRAY)
        image = TF.resize(image, IMAGE_SIZE)
        image = image.float()
        image = image / 255
        label_mask = TF.resize(label_mask, IMAGE_SIZE)
        label_mask = TF.rgb_to_grayscale(label_mask)
        label_mask = label_mask.float()
        return (image, label_mask)

golf_ds = GolfDataset(IMAGES_DIR, LABELMASKS_DIR)
idx = 5
orthophoto = golf_ds.__getitem__(idx)[0]
label_mask = golf_ds.__getitem__(idx)[1]
print('Ortophoto: ', orthophoto.shape, orthophoto)
print('Label:', label_mask.shape, label_mask)