| from torch.functional import Tensor |
| from general_utils import load_model |
| from torch.utils.data import DataLoader |
| import torch |
| import numpy as np |
|
|
| def denorm(img): |
|
|
| np_input = False |
| if isinstance(img, np.ndarray): |
| img = torch.from_numpy(img) |
| np_input = True |
|
|
| mean = torch.Tensor([0.485, 0.456, 0.406]) |
| std = torch.Tensor([0.229, 0.224, 0.225]) |
|
|
| img_denorm = (img*std[:,None,None]) + mean[:,None,None] |
|
|
| if np_input: |
| img_denorm = np.clip(img_denorm.numpy(), 0, 1) |
| else: |
| img_denorm = torch.clamp(img_denorm, 0, 1) |
|
|
| return img_denorm |
|
|
|
|
| def norm(img): |
| mean = torch.Tensor([0.485, 0.456, 0.406]) |
| std = torch.Tensor([0.229, 0.224, 0.225]) |
| return (img - mean[:,None,None]) / std[:,None,None] |
|
|
|
|
| def fast_iou_curve(p, g): |
| |
| g = g[p.sort().indices] |
| p = torch.sigmoid(p.sort().values) |
| |
| scores = [] |
| vals = np.linspace(0, 1, 50) |
|
|
| for q in vals: |
|
|
| n = int(len(g) * q) |
|
|
| valid = torch.where(p > q)[0] |
| if len(valid) > 0: |
| n = int(valid[0]) |
| else: |
| n = len(g) |
|
|
| fn = g[:n].sum() |
| tn = n - fn |
| tp = g[n:].sum() |
| fp = len(g) - n - tp |
|
|
| iou = tp / (tp + fn + fp) |
|
|
| precision = tp / (tp + fp) |
| recall = tp / (tp + fn) |
|
|
| scores += [iou] |
| |
| return vals, scores |
|
|
|
|
| def fast_rp_curve(p, g): |
| |
| g = g[p.sort().indices] |
| p = torch.sigmoid(p.sort().values) |
| |
| precisions, recalls = [], [] |
| vals = np.linspace(p.min(), p.max(), 250) |
|
|
| for q in p[::100000]: |
|
|
| n = int(len(g) * q) |
|
|
| valid = torch.where(p > q)[0] |
| if len(valid) > 0: |
| n = int(valid[0]) |
| else: |
| n = len(g) |
|
|
| fn = g[:n].sum() |
| tn = n - fn |
| tp = g[n:].sum() |
| fp = len(g) - n - tp |
|
|
| iou = tp / (tp + fn + fp) |
|
|
| precision = tp / (tp + fp) |
| recall = tp / (tp + fn) |
|
|
| precisions += [precision] |
| recalls += [recall] |
| |
| return recalls, precisions |
|
|
|
|
| |
|
|
| def img_preprocess(batch, blur=0, grayscale=False, center_context=None, rect=False, rect_color=(255,0,0), rect_width=2, |
| brightness=1.0, bg_fac=1, colorize=False, outline=False, image_size=224): |
| import cv2 |
|
|
| rw = rect_width |
|
|
| out = [] |
| for img, mask in zip(batch[1], batch[2]): |
|
|
| img = img.cpu() if isinstance(img, torch.Tensor) else torch.from_numpy(img) |
| mask = mask.cpu() if isinstance(mask, torch.Tensor) else torch.from_numpy(mask) |
| |
| img *= brightness |
| img_bl = img |
| if blur > 0: |
| img_bl = torch.from_numpy(cv2.GaussianBlur(img.permute(1,2,0).numpy(), (15, 15), blur)).permute(2,0,1) |
| |
| if grayscale: |
| img_bl = img_bl[1][None] |
| |
| |
| |
| img_inp = img*mask + (bg_fac) * img_bl * (1-mask) |
|
|
| if rect: |
| _, bbox = crop_mask(img, mask, context=0.1) |
| img_inp[:, bbox[2]: bbox[3], max(0, bbox[0]-rw):bbox[0]+rw] = torch.tensor(rect_color)[:,None,None] |
| img_inp[:, bbox[2]: bbox[3], max(0, bbox[1]-rw):bbox[1]+rw] = torch.tensor(rect_color)[:,None,None] |
| img_inp[:, max(0, bbox[2]-1): bbox[2]+rw, bbox[0]:bbox[1]] = torch.tensor(rect_color)[:,None,None] |
| img_inp[:, max(0, bbox[3]-1): bbox[3]+rw, bbox[0]:bbox[1]] = torch.tensor(rect_color)[:,None,None] |
|
|
|
|
| if center_context is not None: |
| img_inp = object_crop(img_inp, mask, context=center_context, image_size=image_size) |
|
|
| if colorize: |
| img_gray = denorm(img) |
| img_gray = cv2.cvtColor(img_gray.permute(1,2,0).numpy(), cv2.COLOR_RGB2GRAY) |
| img_gray = torch.stack([torch.from_numpy(img_gray)]*3) |
| img_inp = torch.tensor([1,0.2,0.2])[:,None,None] * img_gray * mask + bg_fac * img_gray * (1-mask) |
| img_inp = norm(img_inp) |
|
|
| if outline: |
| cont = cv2.findContours(mask.byte().numpy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) |
| outline_img = np.zeros(mask.shape, dtype=np.uint8) |
| cv2.drawContours(outline_img, cont[0], -1, thickness=5, color=(255, 255, 255)) |
| outline_img = torch.stack([torch.from_numpy(outline_img)]*3).float() / 255. |
| img_inp = torch.tensor([1,0,0])[:,None,None] * outline_img + denorm(img_inp) * (1- outline_img) |
| img_inp = norm(img_inp) |
|
|
| out += [img_inp] |
|
|
| return torch.stack(out) |
|
|
|
|
| def object_crop(img, mask, context=0.0, square=False, image_size=224): |
| img_crop, bbox = crop_mask(img, mask, context=context, square=square) |
| img_crop = pad_to_square(img_crop, channel_dim=0) |
| img_crop = torch.nn.functional.interpolate(img_crop.unsqueeze(0), (image_size, image_size)).squeeze(0) |
| return img_crop |
| |
|
|
| def crop_mask(img, mask, context=0.0, square=False): |
| |
| assert img.shape[1:] == mask.shape |
| |
| bbox = [mask.max(0).values.argmax(), mask.size(0) - mask.max(0).values.flip(0).argmax()] |
| bbox += [mask.max(1).values.argmax(), mask.size(1) - mask.max(1).values.flip(0).argmax()] |
| bbox = [int(x) for x in bbox] |
| |
| width, height = (bbox[3] - bbox[2]), (bbox[1] - bbox[0]) |
|
|
| |
| if square: |
| bbox[0] = int(max(0, bbox[0] - context * height)) |
| bbox[1] = int(min(mask.size(0), bbox[1] + context * height)) |
| bbox[2] = int(max(0, bbox[2] - context * width)) |
| bbox[3] = int(min(mask.size(1), bbox[3] + context * width)) |
|
|
| width, height = (bbox[3] - bbox[2]), (bbox[1] - bbox[0]) |
| if height > width: |
| bbox[2] = int(max(0, (bbox[2] - 0.5*height))) |
| bbox[3] = bbox[2] + height |
| else: |
| bbox[0] = int(max(0, (bbox[0] - 0.5*width))) |
| bbox[1] = bbox[0] + width |
| else: |
| bbox[0] = int(max(0, bbox[0] - context * height)) |
| bbox[1] = int(min(mask.size(0), bbox[1] + context * height)) |
| bbox[2] = int(max(0, bbox[2] - context * width)) |
| bbox[3] = int(min(mask.size(1), bbox[3] + context * width)) |
|
|
| width, height = (bbox[3] - bbox[2]), (bbox[1] - bbox[0]) |
| img_crop = img[:, bbox[2]: bbox[3], bbox[0]: bbox[1]] |
| return img_crop, bbox |
|
|
|
|
| def pad_to_square(img, channel_dim=2, fill=0): |
| """ |
| |
| |
| add padding such that a squared image is returned """ |
| |
| from torchvision.transforms.functional import pad |
|
|
| if channel_dim == 2: |
| img = img.permute(2, 0, 1) |
| elif channel_dim == 0: |
| pass |
| else: |
| raise ValueError('invalid channel_dim') |
|
|
| h, w = img.shape[1:] |
| pady1 = pady2 = padx1 = padx2 = 0 |
|
|
| if h > w: |
| padx1 = (h - w) // 2 |
| padx2 = h - w - padx1 |
| elif w > h: |
| pady1 = (w - h) // 2 |
| pady2 = w - h - pady1 |
|
|
| img_padded = pad(img, padding=(padx1, pady1, padx2, pady2), padding_mode='constant') |
|
|
| if channel_dim == 2: |
| img_padded = img_padded.permute(1, 2, 0) |
|
|
| return img_padded |
|
|
|
|
| |
|
|
| def split_sentence(inp, limit=9): |
| t_new, current_len = [], 0 |
| for k, t in enumerate(inp.split(' ')): |
| current_len += len(t) + 1 |
| t_new += [t+' '] |
| |
| if current_len > limit and k != len(inp.split(' ')) - 1: |
| current_len = 0 |
| t_new += ['\n'] |
|
|
| t_new = ''.join(t_new) |
| return t_new |
| |
|
|
| from matplotlib import pyplot as plt |
|
|
|
|
| def plot(imgs, *preds, labels=None, scale=1, cmap=plt.cm.magma, aps=None, gt_labels=None, vmax=None): |
| |
| row_off = 0 if labels is None else 1 |
| _, ax = plt.subplots(len(imgs) + row_off, 1 + len(preds), figsize=(scale * float(1 + 2*len(preds)), scale * float(len(imgs)*2))) |
| [a.axis('off') for a in ax.flatten()] |
| |
| if labels is not None: |
| for j in range(len(labels)): |
| t_new = split_sentence(labels[j], limit=6) |
| ax[0, 1+ j].text(0.5, 0.1, t_new, ha='center', fontsize=3+ 10*scale) |
|
|
|
|
| for i in range(len(imgs)): |
| ax[i + row_off,0].imshow(imgs[i]) |
| for j in range(len(preds)): |
| img = preds[j][i][0].detach().cpu().numpy() |
|
|
| if gt_labels is not None and labels[j] == gt_labels[i]: |
| print(j, labels[j], gt_labels[i]) |
| edgecolor = 'red' |
| if aps is not None: |
| ax[i + row_off, 1 + j].text(30, 70, f'AP: {aps[i]:.3f}', color='red', fontsize=8) |
| else: |
| edgecolor = 'k' |
|
|
| rect = plt.Rectangle([0,0], img.shape[0], img.shape[1], facecolor="none", |
| edgecolor=edgecolor, linewidth=3) |
| ax[i + row_off,1 + j].add_patch(rect) |
|
|
| if vmax is None: |
| this_vmax = 1 |
| elif vmax == 'per_prompt': |
| this_vmax = max([preds[j][_i][0].max() for _i in range(len(imgs))]) |
| elif vmax == 'per_image': |
| this_vmax = max([preds[_j][i][0].max() for _j in range(len(preds))]) |
|
|
| ax[i + row_off,1 + j].imshow(img, vmin=0, vmax=this_vmax, cmap=cmap) |
|
|
| |
| |
| plt.tight_layout() |
| plt.subplots_adjust(wspace=0.05, hspace=0.05) |
