random crop and resize while preserving aspect ratio

img_aug_seg.py

import numpy as np
import random

def get_resize_shape(H, W, length=720):
    if H > W and H > length:
        new_H, new_W = length, int(W / H * length)
    elif W > H and W > length:
        new_H, new_W = int(H / W * length), length
    else:
        new_H, new_W = H, W
        
    return new_H, new_W

def get_padding(H, W, crop_size=512):
    lp, rp, up, dp = 0, 0, 0, 0
    if H < crop_size:
        up = (H - crop_size) // 2
        dp = H - crop_size - up
    if W < crop_size:
        lp = (W - crop_size) // 2
        rp = W - crop_size - lp
    return [up, dp, lp, rp]

def resize_and_pad_random(I, mask, crop_size, resize_length):
    Hnew, Wnew = get_resize_shape(I.shape[0], I.shape[1], length=720)
    I2 = cv2.resize(I, (Wnew, Hnew))
    mask2 = cv2.resize(mask, (Wnew, Hnew), interpolation=cv2.INTER_NEAREST)

    [up, dp, lp, rp] = get_padding(Hnew, Wnew)
    [up, dp, lp, rp] = [ x if x >=0 else -x for x in [up, dp, lp, rp]]

    mode = 'constant'
    if(Hnew<Wnew):
        mode = random.sample(['constant', 'edge', 'reflect'], 1)[0]


    colors = random.sample([[(0, 0), (0, 0), (0, 0)], 
                            [(255, 255), (255, 255), (0, 0)],  
                            np.random.randint(0, 255, (3, 2))
                            ], 1)[0]
    I3 = None
    mask3 = None
    if mode == 'constant':
        I3 = np.pad(I2, ((up, dp), (lp, rp), (0, 0)), mode=mode, constant_values=colors)
        mask3 = np.pad(mask2, ((up, dp), (lp, rp)), mode=mode, constant_values=0)
    else:
        I3 = np.pad(I2, ((up, dp), (lp, rp), (0, 0)), mode=mode)
        mask3 = np.pad(mask2, ((up, dp), (lp, rp)), mode=mode)

    Hnew, Wnew = I3.shape[0], I3.shape[1]
    x = 0 if Wnew == crop_size else np.random.randint(0, Wnew - crop_size)
    y = 0 if Hnew == crop_size else np.random.randint(0, Hnew - crop_size)
    I4 = I3[y:y+crop_size, x:x+crop_size]
    mask4 = mask3[y:y+crop_size, x:x+crop_size]
    return I4, mask4