common/utils/cv.py

import math
from typing import Any, List, Union, Tuple, Dict, Optional
import os
import random
from enum import Enum

import cv2
import numpy as np
import pycocotools.mask as maskUtils
from PIL import Image


from .io_utils import json2dict, dict2json


def bbox_intersection(xyxy, xyxy2):
    x1, y1, x2, y2 = xyxy2
    dx1, dy1, dx2, dy2 = xyxy
    ix1, ix2 = max(x1, dx1), min(x2, dx2)
    iy1, iy2 = max(y1, dy1), min(y2, dy2)
    if ix2 >= ix1 and iy2 >= iy1:
        return [ix1, iy1, ix2, iy2]
    return None


def recreate_image(codebook, labels, w, h):
    """Recreate the (compressed) image from the code book & labels"""
    return (codebook[labels].reshape(w, h, -1) * 255).astype(np.uint8)


def quantize_image(image: np.ndarray, n_colors: int, method='kmeans', mask=None):

    from sklearn.cluster import KMeans
    from sklearn.metrics import pairwise_distances_argmin
    from sklearn.utils import shuffle

    # https://scikit-learn.org/stable/auto_examples/cluster/plot_color_quantization.html
    image = np.array(image, dtype=np.float64) / 255

    if len(image.shape) == 3:
        w, h, d = tuple(image.shape)
    else:
        w, h = image.shape
        d = 1

    # assert d == 3
    image_array = image.reshape(-1, d)

    if method == 'kmeans':

        image_array_sample = None
        if mask is not None and not np.all(mask):
            ids  = np.where(mask)
            if len(ids[0]) > 10:
                bg = image[ids][::2]
                fg = image[np.where(mask == 0)]
                max_bg_num = int(fg.shape[0] * 1.5)
                if bg.shape[0] > max_bg_num:
                    bg = shuffle(bg, random_state=0, n_samples=max_bg_num)
                image_array_sample = np.concatenate((fg, bg), axis=0)
                if image_array_sample.shape[0] > 2048:
                    image_array_sample = shuffle(image_array_sample, random_state=0, n_samples=2048)
                else:
                    image_array_sample = None

        if image_array_sample is None:
            image_array_sample = shuffle(image_array, random_state=0, n_samples=2048)
        
        kmeans = KMeans(n_clusters=n_colors, n_init=10, random_state=0).fit(
            image_array_sample
        )

        labels = kmeans.predict(image_array)
        quantized  = recreate_image(kmeans.cluster_centers_, labels, w, h)
        return quantized, kmeans.cluster_centers_, labels

    else:

        codebook_random = shuffle(image_array, random_state=0, n_samples=n_colors)
        labels_random = pairwise_distances_argmin(codebook_random, image_array, axis=0)

        return [recreate_image(codebook_random, labels_random, w, h)]



def get_template_histvq(template: np.ndarray) -> Tuple[List[np.ndarray]]:
    len_shape = len(template.shape)
    num_c = 3
    mask = None
    if len_shape == 2:
        num_c = 1
    elif len_shape == 3 and template.shape[-1] == 4:
        mask = np.where(template[..., -1])
        template = template[..., :num_c][mask]

    values, quantiles = [], []
    for ii in range(num_c):
        v, c = np.unique(template[..., ii].ravel(), return_counts=True)
        q = np.cumsum(c).astype(np.float64)
        if len(q) < 1:
            return None, None
        q /= q[-1]
        values.append(v)
        quantiles.append(q)
    return values, quantiles


def inplace_hist_matching(img: np.ndarray, tv: List[np.ndarray], tq: List[np.ndarray]) -> None:
    len_shape = len(img.shape)
    num_c = 3
    mask = None

    tgtimg = img
    if len_shape == 2:
        num_c = 1
    elif len_shape == 3 and img.shape[-1] == 4:
        mask = np.where(img[..., -1])
        tgtimg = img[..., :num_c][mask]

    im_h, im_w = img.shape[:2]
    oldtype = img.dtype
    for ii in range(num_c):
        _, bin_idx, s_counts = np.unique(tgtimg[..., ii].ravel(), return_inverse=True,
                                                return_counts=True)
        s_quantiles = np.cumsum(s_counts).astype(np.float64)
        if len(s_quantiles) == 0:
            return
        s_quantiles /= s_quantiles[-1]
        interp_t_values = np.interp(s_quantiles, tq[ii], tv[ii]).astype(oldtype)
        if mask is not None:
            img[..., ii][mask] = interp_t_values[bin_idx]
        else:
            img[..., ii] = interp_t_values[bin_idx].reshape((im_h, im_w))
            # try:
            #     img[..., ii] = interp_t_values[bin_idx].reshape((im_h, im_w))
            # except:
            #     LOGGER.error('##################### sth goes wrong')
            #     cv2.imshow('img', img)
            #     cv2.waitKey(0)


def fgbg_hist_matching(fg_list: List, bg: np.ndarray, min_tq_num=128, fg_only=False):
    '''
    inplace op
    '''
    btv, btq = get_template_histvq(bg)
    ftv, ftq = get_template_histvq(fg_list[0])
    num_fg = len(fg_list)
    idx_matched = -1
    if num_fg > 1:
        _ftv, _ftq = get_template_histvq(fg_list[0])
        if _ftq is not None and ftq is not None:
            if len(_ftq[0]) > len(ftq[0]):
                idx_matched = num_fg - 1
                ftv, ftq = _ftv, _ftq
            else:
                idx_matched = 0

    if fg_only and ftq is not None:
        tv, tq = ftv, ftq
        if len(tq[0]) > min_tq_num:
            inplace_hist_matching(bg, tv, tq)
        return

    if btq is not None and ftq is not None:
        if len(btq[0]) > len(ftq[0]):
            tv, tq = btv, btq
            idx_matched = -1
        else:
            tv, tq = ftv, ftq
            if len(tq[0]) > min_tq_num:
                inplace_hist_matching(bg, tv, tq)
        
        if len(tq[0]) > min_tq_num:
            for ii, fg in enumerate(fg_list):
                if ii != idx_matched and len(tq[0]) > min_tq_num:
                    inplace_hist_matching(fg, tv, tq)


def mask2rle(mask: np.ndarray, decode_for_json: bool = True) -> Dict:
    assert mask.ndim == 2
    mask_rle = maskUtils.encode(np.array(
                        mask[..., np.newaxis] > 0, order='F',
                        dtype='uint8'))[0]
    if decode_for_json:
        mask_rle['counts'] = mask_rle['counts'].decode()
    return mask_rle


def rle2mask(rle: Union[Dict, str], to_bool=True):
    # if isinstance(rle, Dict):
    #     rle = rle['counts']
    mask = maskUtils.decode(rle)
    if to_bool:
        return mask > 0
    return mask

def batch_save_masks(masks: np.ndarray, savep: str, compress=None, mask_meta_list=None):
    if isinstance(masks, np.ndarray):
        if masks.ndim == 2:
            masks = masks[None]
    masks = [mask2rle(m) for m in masks]
    if mask_meta_list is not None:
        assert len(masks) == len(mask_meta_list)
        for m, meta in zip(masks, mask_meta_list):
            m.update(meta)
    dict2json(masks, savep, compress=compress)


def batch_load_masks(p: str, to_bool=True):
    masks = json2dict(p)
    masks = [rle2mask(m, to_bool=to_bool) for m in masks]

    return masks


def smart_resize(src: np.ndarray, target_size, upscale_interpolation=cv2.INTER_LINEAR, downscale_interpolation=cv2.INTER_AREA):
    h, w = src.shape[:2]
    th, tw = target_size
    if th == h and tw == w:
        return src.copy()
    if th * tw < h * w:
        interpolation = downscale_interpolation
    else:
        interpolation = upscale_interpolation
    return cv2.resize(src, (tw, th), interpolation=interpolation)



def validate_resolution(resolution: Union[str, Tuple, int], div=-1) -> List:

    '''
    make sure resolution is a valid (h: int, w: int) format and can be divided by div
    '''

    if isinstance(resolution, str):
        resolution = [int(r.strip()) for r in resolution.split(',')]
    elif isinstance(resolution, str):
        resolution = (resolution, resolution)
    elif isinstance(resolution, int):
        resolution = [resolution, resolution]
    elif isinstance(resolution, (Tuple, List)):
        resolution = list(resolution)[:2]

    reso_out = []
    for res in resolution:
        if div > 0 and res % div != 0:
            res = math.ceil(res / div) * div
        reso_out.append(res)
    
    return reso_out


def center_square_pad_resize(img: np.ndarray, target_size, pad_value=0, upscale_interpolation=cv2.INTER_LINEAR, downscale_interpolation=cv2.INTER_AREA, return_pad_info=False):
    h, w = img.shape[:2]
    pad_size = (w, h)
    pad_pos = (0, 0)
    if h != w:
        sz = max(h, w)
        px1 = (sz - w) // 2
        py1 = (sz - h) // 2
        shape = (sz, sz) if img.ndim == 2 else (sz, sz, img.shape[-1])
        padded = np.full(shape, pad_value, dtype=img.dtype)
        padded[py1: py1 + h, px1: px1 + w] = img
        h, w = padded.shape[:2]
        img = padded
        pad_size = (w, h)
        pad_pos = (px1, py1)
    if h != target_size or w != target_size:
        img = smart_resize(img, (target_size, target_size), upscale_interpolation=upscale_interpolation, downscale_interpolation=downscale_interpolation)
    if return_pad_info:
        return img, pad_size, pad_pos
    else:
        return img


def random_hsv(img, hgain: float = 0.015, sgain: float = 0.6, vgain: float = 0.4):
    if hgain or sgain or vgain:
        dtype = img.dtype  # uint8

        r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain]  # random gains
        x = np.arange(0, 256, dtype=r.dtype)
        # lut_hue = ((x * (r[0] + 1)) % 180).astype(dtype)   # original hue implementation from ultralytics<=8.3.78
        lut_hue = ((x + r[0] * 180) % 180).astype(dtype)
        lut_sat = np.clip(x * (r[1] + 1), 0, 255).astype(dtype)
        lut_val = np.clip(x * (r[2] + 1), 0, 255).astype(dtype)
        lut_sat[0] = 0  # prevent pure white changing color, introduced in 8.3.79

        hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_RGB2HSV))
        im_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
        cv2.cvtColor(im_hsv, cv2.COLOR_HSV2RGB, dst=img)  # no return needed

    return img


def resize_short_side_to(src: np.ndarray, short_side: int):
    h, w = src.shape[:2]
    th, tw = h, w
    if h > w:
        tw = short_side
        th = int(round(h / w * short_side))
    else:
        th = short_side
        tw = int(round(w / h * short_side))
    return smart_resize(src, (th, tw))


def random_crop(src: np.ndarray, target_size) -> None:
    '''
    target_size: (h, w)
    '''
    h, w = src.shape[:2]
    th, tw = target_size

    scale = max(th / h, tw / w)
    if scale > 1:
        h = math.ceil(h * scale)
        w = math.ceil(w * scale)
        src = cv2.resize(src, (w, h), interpolation=cv2.INTER_LINEAR)
    
    x0 = y0 = 0
    if h > th:
        y0 = random.randint(0, h - th)
    if w > tw:
        x0 = random.randint(0, w - tw)

    return src[y0: y0 + th, x0: x0 + tw].copy()


def img_bbox(src: np.ndarray):
    if isinstance(src, Image.Image):
        src = np.array(src)
    if src.ndim == 3 and src.shape[-1] == 4:
        src = src[..., 0]
    return cv2.boundingRect(cv2.findNonZero(src.astype(np.uint8)))


def mask_xyxy(mask):
    if isinstance(mask, Image.Image):
        mask = np.array(mask)
    bbox = cv2.boundingRect(cv2.findNonZero(mask.astype(np.uint8)))
    x1, y1, x2, y2 = bbox
    x2 += x1
    y2 += y1
    return [x1, y1, x2, y2]


def argb2rgba(argb):
    return np.concatenate([argb[..., 1:], argb[..., [0]]], axis=2)


def img_alpha_blending(
    drawables: List[np.ndarray], 
    xyxy=None, 
    output_type='numpy', 
    final_size=None, 
    max_depth_val=255,
    premultiplied=True,
):
    '''
    final_size: (h, w)
    '''

    if isinstance(drawables, (np.ndarray, dict)):
        drawables = [drawables]

    # infer final scene size
    if xyxy is not None:
        final_size = [xyxy[3] - xyxy[1], xyxy[2] - xyxy[0]]
        x1, y1, x2, y2 = xyxy
    elif final_size is None:
        d = drawables[0]
        if isinstance(d, dict):
            d = d['img']
        final_size = d.shape[:2]

    final_rgb = np.zeros((final_size[0], final_size[1], 3), dtype=np.float32)
    final_alpha = np.zeros_like(final_rgb[..., [0]])
    final_depth = None

    for drawable_img in drawables:
        dxyxy = None
        depth = None
        if isinstance(drawable_img, dict):
            depth = drawable_img.get('depth', None)
            tag = drawable_img.get('tag', None)
            if depth is not None:
                if depth.ndim == 2:
                    depth = depth[..., None]
                if final_depth is None:
                    final_depth = np.full_like(final_alpha, fill_value=max_depth_val)
            if 'xyxy' in drawable_img:
                dxyxy = drawable_img['xyxy']
                dx1, dy1, dx2, dy2 = dxyxy
            drawable_img = drawable_img['img']
            if dxyxy is not None:
                if dx1 < 0:
                    drawable_img = drawable_img[:, -dx1:]
                    if depth is not None:
                        depth = depth[:, -dx1:]
                    dx1 = 0
                if dy1 < 0:
                    drawable_img = drawable_img[-dy1:]
                    if depth is not None:
                        depth = depth[-dy1:]
                    dy1 = 0

        if drawable_img.ndim == 3 and drawable_img.shape[-1] == 3:
            drawable_alpha = np.ones_like(drawable_img[..., [-1]])
        else:
            drawable_alpha = drawable_img[..., [-1]] / 255

        drawable_img = drawable_img[..., :3]

        if xyxy is not None:
            if dxyxy is None:
                drawable_img = drawable_img[y1: y2, x1: x2]
            else:
                intersection = bbox_intersection(xyxy, dxyxy)
                if intersection is None:
                    continue
                ix1, iy1, ix2, iy2 = intersection
                if depth is not None:
                    depth = depth[iy1-dy1: iy2-dy1, ix1-dx1: ix2-dx1]
                    drawable_alpha = drawable_alpha[iy1-dy1: iy2-dy1, ix1-dx1: ix2-dx1]
                    update_mask = (final_depth[iy1-y1: iy2-y1, ix1-x1: ix2-x1] > depth).astype(np.uint8)
                    final_depth[iy1-y1: iy2-y1, ix1-x1: ix2-x1] = update_mask * depth + (1-update_mask) * final_depth[iy1-y1: iy2-y1, ix1-x1: ix2-x1]
                    drawable_img = drawable_img[iy1-dy1: iy2-dy1, ix1-dx1: ix2-dx1]
                    final_rgb[iy1-y1: iy2-y1, ix1-x1: ix2-x1] = update_mask * (final_rgb[iy1-y1: iy2-y1, ix1-x1: ix2-x1] * (1-drawable_alpha) + drawable_img) + \
                        (1 - update_mask) * (drawable_img * (1-final_alpha[iy1-y1: iy2-y1, ix1-x1: ix2-x1]) + final_rgb[iy1-y1: iy2-y1, ix1-x1: ix2-x1])
                    final_alpha[iy1-y1: iy2-y1, ix1-x1: ix2-x1] = np.clip(final_alpha[iy1-y1: iy2-y1, ix1-x1: ix2-x1] + drawable_alpha, 0, 1)
                else:
                    
                    drawable_alpha = drawable_alpha[iy1-dy1: iy2-dy1, ix1-dx1: ix2-dx1]
                    final_alpha[iy1-y1: iy2-y1, ix1-x1: ix2-x1] += drawable_alpha
                    drawable_img = drawable_img[iy1-dy1: iy2-dy1, ix1-dx1: ix2-dx1]
                    final_rgb[iy1-y1: iy2-y1, ix1-x1: ix2-x1] = final_rgb[iy1-y1: iy2-y1, ix1-x1: ix2-x1] * (1-drawable_alpha) + drawable_img
                    continue

        elif dxyxy is None:
            if depth is not None:
                update_mask = (final_depth > depth).astype(np.uint8)
                final_depth = update_mask * depth + (1-update_mask) * final_depth
                final_rgb = update_mask * (final_rgb * (1-drawable_alpha) + drawable_img) + \
                    (1 - update_mask) * (drawable_img * (1-final_alpha) + final_rgb)
                final_alpha = np.clip(final_alpha + drawable_alpha, 0, 1)
            else:
                final_alpha += drawable_alpha
                final_alpha = np.clip(final_alpha, 0, 1)
                if not premultiplied:
                    drawable_img = drawable_img * drawable_alpha
                final_rgb = final_rgb * (1 - drawable_alpha) + drawable_img
        else:
            if depth is not None:
                update_mask = (final_depth[dy1: dy2, dx1: dx2] > depth).astype(np.uint8)
                update_mask = update_mask * (drawable_alpha > 0.1)
                final_depth[dy1: dy2, dx1: dx2] = update_mask * depth + (1-update_mask) * final_depth[dy1: dy2, dx1: dx2]
                final_rgb[dy1: dy2, dx1: dx2] = update_mask * (final_rgb[dy1: dy2, dx1: dx2] * (1-drawable_alpha) + drawable_img) + \
                    (1 - update_mask) * (drawable_img * (1-final_alpha[dy1: dy2, dx1: dx2]) + final_rgb[dy1: dy2, dx1: dx2])
                final_alpha[dy1: dy2, dx1: dx2] = np.clip(final_alpha[dy1: dy2, dx1: dx2] + drawable_alpha, 0, 1)
            else:
                final_alpha[dy1: dy2, dx1: dx2] += drawable_alpha
                final_alpha = np.clip(final_alpha, 0, 1)
                if not premultiplied:
                    drawable_img = drawable_img * drawable_alpha
                final_rgb[dy1: dy2, dx1: dx2] = final_rgb[dy1: dy2, dx1: dx2] * (1-drawable_alpha) + drawable_img

    final_alpha = np.clip(final_alpha, 0, 1) * 255
    final = np.concatenate([final_rgb, final_alpha], axis=2)
    final = np.clip(final, 0, 255).astype(np.uint8)

    output_type = output_type.lower()
    if output_type == 'pil':
        final = Image.fromarray(final)
    elif output_type == 'dict':
        final = {
            'img': final
        }
        if final_depth is not None:
            final['depth'] = final_depth

    return final


def rgba_to_rgb_fixbg(img: np.ndarray, background_color=255):
    if isinstance(img, Image.Image):
        img = np.array(img)
    assert img.ndim == 3
    if img.shape[-1] == 3:
        return img
    if isinstance(background_color, int):
        bg = np.full_like(img[..., :3], fill_value=background_color)
    else:
        background_color = np.array(background_color)[:3].astype(np.uint8)
        bg = np.full_like(img[..., :3], fill_value=255)
        bg[..., :3] = background_color
    return img_alpha_blending([bg, img])[..., :3].copy()


def build_alpha_pyramid(color, alpha, dk=1.2):
    # Written by lvmin at Stanford
    # Massive iterative Gaussian filters are mathematically consistent to pyramid.

    pyramid = []
    current_premultiplied_color = color * alpha
    current_alpha = alpha

    while True:
        pyramid.append((current_premultiplied_color, current_alpha))

        H, W, C = current_alpha.shape
        if min(H, W) == 1:
            break

        current_premultiplied_color = cv2.resize(current_premultiplied_color, (int(W / dk), int(H / dk)), interpolation=cv2.INTER_AREA)
        current_alpha = cv2.resize(current_alpha, (int(W / dk), int(H / dk)), interpolation=cv2.INTER_AREA)[:, :, None]
    return pyramid[::-1]


def pad_rgb(np_rgba_hwc_uint8, return_format='rgb', to_uint8=False, keep_ori_pixel=True):
    # Written by lvmin at Stanford
    # Massive iterative Gaussian filters are mathematically consistent to pyramid.


    np_rgba_hwc = np_rgba_hwc_uint8.astype(np.float32) / 255.0

    if keep_ori_pixel:
        ori_rgb = np_rgba_hwc[..., :3].copy()
        ori_alpha = np_rgba_hwc[..., [-1]].copy()

    pyramid = build_alpha_pyramid(color=np_rgba_hwc[..., :3], alpha=np_rgba_hwc[..., 3:])

    top_c, top_a = pyramid[0]
    fg = np.sum(top_c, axis=(0, 1), keepdims=True) / np.sum(top_a, axis=(0, 1), keepdims=True).clip(1e-8, 1e32)

    for layer_c, layer_a in pyramid:
        layer_h, layer_w, _ = layer_c.shape
        fg = cv2.resize(fg, (layer_w, layer_h), interpolation=cv2.INTER_LINEAR)
        fg = layer_c + fg * (1.0 - layer_a)

    if keep_ori_pixel:
        fg = np.clip(ori_alpha * ori_rgb + (1-ori_alpha) * fg, 0, 1)

    if return_format == 'argb':
        fg = np.concatenate([np_rgba_hwc[..., 3:], fg], axis=2)

    if to_uint8:
        fg = (fg * 255).astype(np.uint8)

    return fg


def checkerboard_vis(img):
    # y = y.clip(0, 1).movedim(1, -1)
    # alpha = y[..., :1]
    # fg = y[..., 1:]
    H, W, C = img.shape
    img = img.astype(np.float32) / 255.
    alpha = img[..., [-1]]
    fg = img[..., :3]
    cb = checkerboard(shape=(H // 64, W // 64))
    cb = cv2.resize(cb, (W, H), interpolation=cv2.INTER_NEAREST)
    cb = (0.5 + (cb - 0.5) * 0.1)[..., None]
    # cb = torch.from_numpy(cb).to(fg)

    vis = fg * alpha + cb * (1 - alpha)
    vis = (vis * 255.0).clip(0, 255).astype(np.uint8)
    return vis


def checkerboard(shape):
    return np.indices(shape).sum(axis=0) % 2


def visualize_rgba(rgba):
    rgba = rgba.astype(np.float32) / 255.
    H, W, C = rgba.shape
    cb = checkerboard(shape=(H // 64, W // 64))
    cb = cv2.resize(cb, (W, H), interpolation=cv2.INTER_NEAREST)
    cb = (0.5 + (cb - 0.5) * 0.1)[..., None]
    alpha = rgba[..., [-1]]
    fg = rgba[..., :3]
    vis = (fg * alpha + cb * (1 - alpha))
    vis = np.clip(vis * 255.0, 0, 255).astype(np.uint8)
    return vis


class DrawMethod(Enum):
    LINE = 'line'
    CIRCLE = 'circle'
    SQUARE = 'square'


def make_random_rectangle_mask(shape, margin=10, bbox_min_size=64, bbox_max_size=384, min_times=1, max_times=4):
    height, width = shape
    mask = np.zeros((height, width), np.float32)
    bbox_max_size = min(bbox_max_size, height - margin * 2, width - margin * 2)
    times = np.random.randint(min_times, max_times + 1)
    for i in range(times):
        box_width = np.random.randint(bbox_min_size, bbox_max_size)
        box_height = np.random.randint(bbox_min_size, bbox_max_size)
        start_x = np.random.randint(margin, width - margin - box_width + 1)
        start_y = np.random.randint(margin, height - margin - box_height + 1)
        mask[start_y:start_y + box_height, start_x:start_x + box_width] = 1
    return mask


def make_random_irregular_mask(shape, max_angle=4, max_len=600, max_width=256, min_times=1, max_times=5,
                               draw_method=DrawMethod.LINE):
    draw_method = DrawMethod(draw_method)

    height, width = shape
    mask = np.zeros((height, width), np.float32)
    times = np.random.randint(min_times, max_times + 1)
    for i in range(times):
        start_x = np.random.randint(width)
        start_y = np.random.randint(height)
        for j in range(1 + np.random.randint(5)):
            angle = 0.01 + np.random.randint(max_angle)
            if i % 2 == 0:
                angle = 2 * 3.1415926 - angle
            length = 10 + np.random.randint(max_len)
            brush_w = 10 + np.random.randint(max_width)
            end_x = np.clip((start_x + length * np.sin(angle)).astype(np.int32), 0, width)
            end_y = np.clip((start_y + length * np.cos(angle)).astype(np.int32), 0, height)
            if draw_method == DrawMethod.LINE:
                cv2.line(mask, (start_x, start_y), (end_x, end_y), 1.0, brush_w)
            elif draw_method == DrawMethod.CIRCLE:
                cv2.circle(mask, (start_x, start_y), radius=brush_w, color=1., thickness=-1)
            elif draw_method == DrawMethod.SQUARE:
                radius = brush_w // 2
                mask[start_y - radius:start_y + radius, start_x - radius:start_x + radius] = 1
            start_x, start_y = end_x, end_y
    return mask


def random_pad_img(img: np.array, tmax=0, bmax=0, lmax=0, rmax=0, pad_values=0):
    l = r = t = b = 0
    if tmax > 0:
        t = random.randint(0, tmax)
    if bmax > 0:
        b = random.randint(0, bmax)
    if lmax > 0:
        l = random.randint(0, lmax)
    if rmax > 0:
        r = random.randint(0, rmax)
    if t > 0 or b > 0 or l > 0 or r > 0:
        padded = cv2.copyMakeBorder(img, t, b, l, r, borderType=cv2.BORDER_CONSTANT, value=pad_values)
    else:
        padded = img.copy()

    return padded, (t, b, l, r)