faisalhr1997/codeformer

1import cv22import random3 4 5def mod_crop(img, scale):6    """Mod crop images, used during testing.7 8    Args:9        img (ndarray): Input image.10        scale (int): Scale factor.11 12    Returns:13        ndarray: Result image.14    """15    img = img.copy()16    if img.ndim in (2, 3):17        h, w = img.shape[0], img.shape[1]18        h_remainder, w_remainder = h % scale, w % scale19        img = img[:h - h_remainder, :w - w_remainder, ...]20    else:21        raise ValueError(f'Wrong img ndim: {img.ndim}.')22    return img23 24 25def paired_random_crop(img_gts, img_lqs, gt_patch_size, scale, gt_path):26    """Paired random crop.27 28    It crops lists of lq and gt images with corresponding locations.29 30    Args:31        img_gts (list[ndarray] | ndarray): GT images. Note that all images32            should have the same shape. If the input is an ndarray, it will33            be transformed to a list containing itself.34        img_lqs (list[ndarray] | ndarray): LQ images. Note that all images35            should have the same shape. If the input is an ndarray, it will36            be transformed to a list containing itself.37        gt_patch_size (int): GT patch size.38        scale (int): Scale factor.39        gt_path (str): Path to ground-truth.40 41    Returns:42        list[ndarray] | ndarray: GT images and LQ images. If returned results43            only have one element, just return ndarray.44    """45 46    if not isinstance(img_gts, list):47        img_gts = [img_gts]48    if not isinstance(img_lqs, list):49        img_lqs = [img_lqs]50 51    h_lq, w_lq, _ = img_lqs[0].shape52    h_gt, w_gt, _ = img_gts[0].shape53    lq_patch_size = gt_patch_size // scale54 55    if h_gt != h_lq * scale or w_gt != w_lq * scale:56        raise ValueError(f'Scale mismatches. GT ({h_gt}, {w_gt}) is not {scale}x ',57                         f'multiplication of LQ ({h_lq}, {w_lq}).')58    if h_lq < lq_patch_size or w_lq < lq_patch_size:59        raise ValueError(f'LQ ({h_lq}, {w_lq}) is smaller than patch size '60                         f'({lq_patch_size}, {lq_patch_size}). '61                         f'Please remove {gt_path}.')62 63    # randomly choose top and left coordinates for lq patch64    top = random.randint(0, h_lq - lq_patch_size)65    left = random.randint(0, w_lq - lq_patch_size)66 67    # crop lq patch68    img_lqs = [v[top:top + lq_patch_size, left:left + lq_patch_size, ...] for v in img_lqs]69 70    # crop corresponding gt patch71    top_gt, left_gt = int(top * scale), int(left * scale)72    img_gts = [v[top_gt:top_gt + gt_patch_size, left_gt:left_gt + gt_patch_size, ...] for v in img_gts]73    if len(img_gts) == 1:74        img_gts = img_gts[0]75    if len(img_lqs) == 1:76        img_lqs = img_lqs[0]77    return img_gts, img_lqs78 79 80def augment(imgs, hflip=True, rotation=True, flows=None, return_status=False):81    """Augment: horizontal flips OR rotate (0, 90, 180, 270 degrees).82 83    We use vertical flip and transpose for rotation implementation.84    All the images in the list use the same augmentation.85 86    Args:87        imgs (list[ndarray] | ndarray): Images to be augmented. If the input88            is an ndarray, it will be transformed to a list.89        hflip (bool): Horizontal flip. Default: True.90        rotation (bool): Ratotation. Default: True.91        flows (list[ndarray]: Flows to be augmented. If the input is an92            ndarray, it will be transformed to a list.93            Dimension is (h, w, 2). Default: None.94        return_status (bool): Return the status of flip and rotation.95            Default: False.96 97    Returns:98        list[ndarray] | ndarray: Augmented images and flows. If returned99            results only have one element, just return ndarray.100 101    """102    hflip = hflip and random.random() < 0.5103    vflip = rotation and random.random() < 0.5104    rot90 = rotation and random.random() < 0.5105 106    def _augment(img):107        if hflip:  # horizontal108            cv2.flip(img, 1, img)109        if vflip:  # vertical110            cv2.flip(img, 0, img)111        if rot90:112            img = img.transpose(1, 0, 2)113        return img114 115    def _augment_flow(flow):116        if hflip:  # horizontal117            cv2.flip(flow, 1, flow)118            flow[:, :, 0] *= -1119        if vflip:  # vertical120            cv2.flip(flow, 0, flow)121            flow[:, :, 1] *= -1122        if rot90:123            flow = flow.transpose(1, 0, 2)124            flow = flow[:, :, [1, 0]]125        return flow126 127    if not isinstance(imgs, list):128        imgs = [imgs]129    imgs = [_augment(img) for img in imgs]130    if len(imgs) == 1:131        imgs = imgs[0]132 133    if flows is not None:134        if not isinstance(flows, list):135            flows = [flows]136        flows = [_augment_flow(flow) for flow in flows]137        if len(flows) == 1:138            flows = flows[0]139        return imgs, flows140    else:141        if return_status:142            return imgs, (hflip, vflip, rot90)143        else:144            return imgs145 146 147def img_rotate(img, angle, center=None, scale=1.0):148    """Rotate image.149 150    Args:151        img (ndarray): Image to be rotated.152        angle (float): Rotation angle in degrees. Positive values mean153            counter-clockwise rotation.154        center (tuple[int]): Rotation center. If the center is None,155            initialize it as the center of the image. Default: None.156        scale (float): Isotropic scale factor. Default: 1.0.157    """158    (h, w) = img.shape[:2]159 160    if center is None:161        center = (w // 2, h // 2)162 163    matrix = cv2.getRotationMatrix2D(center, angle, scale)164    rotated_img = cv2.warpAffine(img, matrix, (w, h))165    return rotated_img166