File size: 5,759 Bytes
9e14838 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 | # -*- coding: utf-8 -*-
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)
# Fix bug RGBA
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)
# Resize mask and apply colormap
overlay = mask.resize(img.size, resample=Image.BICUBIC)
overlay = (255 * cmap(np.asarray(overlay) ** 1)[:, :, :3]).astype(np.uint8)
# Overlay the image with the mask
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)
# to [16/255, 235/255]
img_ycbcr[:, :, 0] = (img_ycbcr[:, :, 0] * (235 - 16) + 16) / 255.0
# to [16/255, 240/255]
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)
# to [0, 1]
img_ycbcr[:, :, 0] = (img_ycbcr[:, :, 0] * 255.0 - 16) / (235 - 16)
# to [0, 1]
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
|