RandomSizedCrop.py

import os
import gc
import cv2
import glob
import shutil
import random
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm


H = 320
W = 480


IMGS = r"aerial_dataset_resized\train\images\*.jpg"
MSKS = r"aerial_dataset_resized\train\masks\*.png"

for img_path, msk_path in zip(tqdm(glob.glob(IMGS)), glob.glob(MSKS)):
    raw_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
    raw_msk = cv2.imread(msk_path, cv2.IMREAD_COLOR)

    raw_h, raw_w, _ = raw_img.shape

    h_limit = raw_h - H
    w_limit = raw_w - W

    h_start = random.randint(0, h_limit)
    w_start = random.randint(0, w_limit)
    # print(w_start, h_start)

    w_end = w_start + W
    h_end = h_start + H

    img_ = raw_img[h_start:h_end, w_start:w_end]
    msk_ = raw_msk[h_start:h_end, w_start:w_end]

    # assert img_.shape == (H, W, 3), f"lol noob, {os.path.split(img_path)[-1]}, {w_end - w_start, h_end - h_start}"

    # plt.subplot(1, 2, 1)
    # plt.imshow(img_)
    # plt.axis("off")

    # plt.subplot(1, 2, 2)
    # plt.imshow(msk_)
    # plt.axis("off")
    # plt.show()
    # break