|
import os |
|
import numpy as np |
|
from PIL import Image |
|
from scipy.ndimage.morphology import binary_dilation |
|
import cv2 |
|
|
|
from numpy.linalg import norm |
|
|
|
import torch |
|
from torch.nn import functional as NF |
|
from torchvision.transforms import functional as TF |
|
|
|
|
|
color_palette = [0, 0, 0, 0, 0, 128, 0, 128, 0, 128, 0, 0] + [100, 100, 100] * 252 |
|
|
|
|
|
def postprocessing_pred(pred: np.array) -> np.array: |
|
|
|
label_cnt, labels = cv2.connectedComponentsWithAlgorithm(pred, 8, cv2.CV_32S, cv2.CCL_GRANA) |
|
if label_cnt == 2: |
|
if labels[0, 0] == pred[0, 0]: |
|
pred = labels |
|
else: |
|
pred = 1 - labels |
|
else: |
|
max_cnt, max_label = 0, 0 |
|
for i in range(label_cnt): |
|
mask = labels == i |
|
if pred[mask][0] == 0: |
|
continue |
|
cnt = len(mask.nonzero()[0]) |
|
if cnt > max_cnt: |
|
max_cnt = cnt |
|
max_label = i |
|
pred = labels == max_label |
|
|
|
return pred.astype(np.uint8) |
|
|
|
|
|
def calc_uncertainty(score): |
|
|
|
|
|
score_top, _ = score.topk(k=2, dim=1) |
|
uncertainty = score_top[:, 0] / (score_top[:, 1] + 1e-8) |
|
uncertainty = torch.exp(1 - uncertainty).unsqueeze(1) |
|
return uncertainty |
|
|
|
|
|
def save_seg_mask(pred, seg_path, palette=color_palette): |
|
|
|
seg_img = Image.fromarray(pred) |
|
seg_img.putpalette(palette) |
|
seg_img.save(seg_path) |
|
|
|
|
|
def add_overlay(img, mask, colors=color_palette, alpha=0.4, cscale=1): |
|
|
|
ids = np.unique(mask) |
|
img_overlay = img.copy() |
|
ones_np = np.ones(img.shape) * (1 - alpha) |
|
|
|
colors = np.reshape(colors, (-1, 3)) |
|
colors = np.atleast_2d(colors) * cscale |
|
|
|
for i in ids[1:]: |
|
|
|
canvas = img * alpha + ones_np * np.array(colors[i])[::-1] |
|
|
|
binary_mask = mask == i |
|
img_overlay[binary_mask] = canvas[binary_mask] |
|
|
|
contour = binary_dilation(binary_mask) ^ binary_mask |
|
img_overlay[contour, :] = 0 |
|
|
|
return img_overlay |
|
|
|
|
|
def save_overlay(img, mask, overlay_path, colors=[255, 0, 0], alpha=0.4, cscale=1): |
|
|
|
img = (img.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8) |
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) |
|
|
|
img_overlay = add_overlay(img, mask, colors, alpha, cscale) |
|
cv2.imwrite(overlay_path, img_overlay) |
|
|
|
|
|
def load_image_in_PIL(path, mode='RGB'): |
|
img = Image.open(path) |
|
img.load() |
|
return img.convert(mode) |
|
|
|
|
|
def normalize(x): |
|
return x / norm(x, ord=2, axis=1, keepdims=True) |
|
|
|
|
|
def dist(p0, p1, axis): |
|
return norm(p0 - p1, ord=2, axis=axis) |
|
|
|
|
|
def resize_img(img, out_size): |
|
h, w = img.shape[:2] |
|
|
|
if h > w: |
|
w_new = int(out_size * w / h) |
|
h_new = out_size |
|
else: |
|
h_new = int(out_size * h / w) |
|
w_new = out_size |
|
|
|
img = cv2.resize(img, (w_new, h_new)) |
|
return img |
|
|
|
|
|
def unify_features(features): |
|
output_size = features['f0'].shape[-2:] |
|
feature_tuple = tuple() |
|
|
|
for key, f in features.items(): |
|
if key != 'f0': |
|
f = NF.interpolate( |
|
f, |
|
size=output_size, mode='bilinear', align_corners=False |
|
) |
|
feature_tuple += (f,) |
|
|
|
unified_feature = torch.cat(feature_tuple, dim=1) |
|
|
|
return unified_feature |
|
|
|
|
|
def pad_divide_by(in_list, d, in_size): |
|
out_list = [] |
|
h, w = in_size |
|
if h % d > 0: |
|
new_h = h + d - h % d |
|
else: |
|
new_h = h |
|
if w % d > 0: |
|
new_w = w + d - w % d |
|
else: |
|
new_w = w |
|
lh, uh = int((new_h - h) / 2), int(new_h - h) - int((new_h - h) / 2) |
|
lw, uw = int((new_w - w) / 2), int(new_w - w) - int((new_w - w) / 2) |
|
pad_array = (int(lw), int(uw), int(lh), int(uh)) |
|
for inp in in_list: |
|
out_list.append(NF.pad(inp, pad_array)) |
|
|
|
return out_list, pad_array |
|
|
