| |
| import math |
| import os |
| import random |
|
|
| import cv2 |
| import numpy as np |
| from matplotlib import cm |
| from PIL import Image |
|
|
|
|
| def load_image(image_path): |
| """Loading image""" |
| img = Image.open(image_path) |
| |
| if img.mode != "RGB": |
| img = img.convert("RGB") |
| img = np.array(img) |
| return img |
|
|
|
|
| def crop_by_margin(image, margin=[0, 0]): |
| """Cropping images by margins as a step of preprocessing""" |
| H, W = image.shape[:2] |
| margin_x, margin_y = margin |
| image = image[margin_y : H - margin_y, margin_x : W - margin_x, :] |
| return image |
|
|
|
|
| def gaussian_radius(det_size, min_overlap=0.7): |
| """Calculating gaussian radius to compute std for Unnormalized Gaussian Mask""" |
| height, width = det_size |
|
|
| a1 = 1 |
| b1 = height + width |
| c1 = width * height * (1 - min_overlap) / (1 + min_overlap) |
| sq1 = np.sqrt(b1**2 - 4 * a1 * c1) |
| r1 = (b1 + sq1) / 2 |
|
|
| a2 = 4 |
| b2 = 2 * (height + width) |
| c2 = (1 - min_overlap) * width * height |
| sq2 = np.sqrt(b2**2 - 4 * a2 * c2) |
| r2 = (b2 + sq2) / 2 |
|
|
| a3 = 4 * min_overlap |
| b3 = -2 * min_overlap * (height + width) |
| c3 = (min_overlap - 1) * width * height |
| sq3 = np.sqrt(b3**2 - 4 * a3 * c3) |
| r3 = (b3 + sq3) / 2 |
| return min(r1, r2, r3) |
|
|
|
|
| def cal_mask_wh(p, mask): |
| """Adaptively calculating blending mask W, H at the most vulnerable points perspective""" |
| cy, cx = p |
| mask_h, mask_w = mask.shape |
| w = 0 |
| h = 0 |
|
|
| for i in [-1, 1]: |
| shift_y = 0 |
| while ( |
| (cy + shift_y > -mask_h) |
| and (cy + shift_y < mask_h) |
| and (mask[cy + shift_y, cx] > 128) |
| ): |
| w += 1 |
| shift_y += i |
|
|
| shift_x = 0 |
| while ( |
| (cx + shift_x > -mask_w) |
| and (cx + shift_x < mask_w) |
| and (mask[cy, cx + shift_x] > 128) |
| ): |
| h += 1 |
| shift_x += i |
|
|
| return w, h |
|
|
|
|
| def overlay_mask( |
| img: Image.Image, mask: Image.Image, colormap: str = "jet", alpha: float = 0.7 |
| ) -> Image.Image: |
| """Overlay a colormapped mask on a background image |
| |
| >>> from PIL import Image |
| >>> import matplotlib.pyplot as plt |
| >>> from torchcam.utils import overlay_mask |
| >>> img = ... |
| >>> cam = ... |
| >>> overlay = overlay_mask(img, cam) |
| |
| Args: |
| img: background image |
| mask: mask to be overlayed in grayscale |
| colormap: colormap to be applied on the mask |
| alpha: transparency of the background image |
| |
| Returns: |
| overlayed image |
| |
| Raises: |
| TypeError: when the arguments have invalid types |
| ValueError: when the alpha argument has an incorrect value |
| """ |
|
|
| if not isinstance(img, Image.Image) or not isinstance(mask, Image.Image): |
| raise TypeError("img and mask arguments need to be PIL.Image") |
|
|
| if not isinstance(alpha, float) or alpha < 0 or alpha >= 1: |
| raise ValueError( |
| "alpha argument is expected to be of type float between 0 and 1" |
| ) |
|
|
| cmap = cm.get_cmap(colormap) |
| |
| overlay = mask.resize(img.size, resample=Image.BICUBIC) |
| overlay = (255 * cmap(np.asarray(overlay) ** 1)[:, :, :3]).astype(np.uint8) |
| |
| overlayed_img = Image.fromarray( |
| (alpha * np.asarray(img) + (1 - alpha) * overlay).astype(np.uint8) |
| ) |
|
|
| return overlayed_img |
|
|
|
|
| def bgr2ycbcr(img_bgr): |
| img_bgr = img_bgr.astype(np.float32) |
| img_ycrcb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2YCR_CB) |
| img_ycbcr = img_ycrcb[:, :, (0, 2, 1)].astype(np.float32) |
| |
| img_ycbcr[:, :, 0] = (img_ycbcr[:, :, 0] * (235 - 16) + 16) / 255.0 |
| |
| img_ycbcr[:, :, 1:] = (img_ycbcr[:, :, 1:] * (240 - 16) + 16) / 255.0 |
|
|
| return img_ycbcr |
|
|
|
|
| def ycbcr2bgr(img_ycbcr): |
| img_ycbcr = img_ycbcr.astype(np.float32) |
| |
| img_ycbcr[:, :, 0] = (img_ycbcr[:, :, 0] * 255.0 - 16) / (235 - 16) |
| |
| img_ycbcr[:, :, 1:] = (img_ycbcr[:, :, 1:] * 255.0 - 16) / (240 - 16) |
| img_ycrcb = img_ycbcr[:, :, (0, 2, 1)].astype(np.float32) |
| img_bgr = cv2.cvtColor(img_ycrcb, cv2.COLOR_YCR_CB2BGR) |
|
|
| return img_bgr |
|
|
|
|
| def gaussian_noise_color(img, param=None): |
| if param is None: |
| param = [0.001, 0.002, 0.005, 0.01, 0.05] |
| ycbcr = bgr2ycbcr(img) / 255 |
| size_a = ycbcr.shape |
| b = ( |
| ycbcr + math.sqrt(param) * np.random.randn(size_a[0], size_a[1], size_a[2]) |
| ) * 255 |
| b = ycbcr2bgr(b) |
| img = np.clip(b, 0, 255).astype(np.uint8) |
| return img |
|
|
|
|
| def block_wise(img, param): |
| width = 8 |
| block = np.ones((width, width, 3)).astype(int) * 128 |
| param = min(img.shape[0], img.shape[1]) // 256 * param |
| for i in range(param): |
| r_w = random.randint(0, img.shape[1] - 1 - width) |
| r_h = random.randint(0, img.shape[0] - 1 - width) |
| img[r_h : r_h + width, r_w : r_w + width, :] = block |
|
|
| return img |
|
|
|
|
| def color_saturation(img, param): |
| ycbcr = bgr2ycbcr(img) |
| ycbcr[:, :, 1] = 0.5 + (ycbcr[:, :, 1] - 0.5) * param |
| ycbcr[:, :, 2] = 0.5 + (ycbcr[:, :, 2] - 0.5) * param |
| img = ycbcr2bgr(ycbcr).astype(np.uint8) |
|
|
| return img |
|
|
|
|
| def color_contrast(img, param): |
| img = img.astype(np.float32) * param |
| img = img.astype(np.uint8) |
|
|
| return img |
|
|
|
|
| def gaussian_blur(img, param): |
| img = cv2.GaussianBlur(img, (param, param), param * 1.0 / 6) |
|
|
| return img |
|
|
|
|
| def jpeg_compression(img, param): |
| h, w, _ = img.shape |
| s_h = h // param |
| s_w = w // param |
| img = cv2.resize(img, (s_w, s_h)) |
| img = cv2.resize(img, (w, h)) |
|
|
| return img |
|
|
|
|
| def video_compression(vid_in, vid_out, param): |
| cmd = f"ffmpeg -i {vid_in} -crf {param} -y {vid_out}" |
| os.system(cmd) |
|
|
| return |
|
|
