Upload utils.py

utils/utils.py

@@ -0,0 +1,140 @@
+import numpy as np
+import cv2
+import math
+
+def f(r, T=0.6, beta=0.1):
+    return np.where(r < T, beta + (1 - beta) / T * r, 1)
+
+# Get the bounding box of the mask
+def get_bbox_from_mask(mask):
+    h,w = mask.shape[0],mask.shape[1]
+
+    if mask.sum() < 10:
+        return 0,h,0,w
+    rows = np.any(mask,axis=1)
+    cols = np.any(mask,axis=0)
+    y1,y2 = np.where(rows)[0][[0,-1]]
+    x1,x2 = np.where(cols)[0][[0,-1]]
+    return (y1,y2,x1,x2)
+
+# Expand the bounding box
+def expand_bbox(mask, yyxx, ratio, min_crop=0):
+    y1,y2,x1,x2 = yyxx
+    H,W = mask.shape[0], mask.shape[1]
+
+    yyxx_area = (y2-y1+1) * (x2-x1+1)
+    r1 = yyxx_area / (H * W)
+    r2 = f(r1)
+    ratio = math.sqrt(r2 / r1)
+
+    xc, yc = 0.5 * (x1 + x2), 0.5 * (y1 + y2)
+    h = ratio * (y2-y1+1)
+    w = ratio * (x2-x1+1)
+    h = max(h,min_crop)
+    w = max(w,min_crop)
+
+    x1 = int(xc - w * 0.5)
+    x2 = int(xc + w * 0.5)
+    y1 = int(yc - h * 0.5)
+    y2 = int(yc + h * 0.5)
+
+    x1 = max(0,x1)
+    x2 = min(W,x2)
+    y1 = max(0,y1)
+    y2 = min(H,y2)
+    return (y1,y2,x1,x2)
+
+# Pad the image to a square shape
+def pad_to_square(image, pad_value = 255, random = False):
+    H,W = image.shape[0], image.shape[1]
+    if H == W:
+        return image
+
+    padd = abs(H - W)
+    if random:
+        padd_1 = int(np.random.randint(0,padd))
+    else:
+        padd_1 = int(padd / 2)
+    padd_2 = padd - padd_1
+
+    if len(image.shape) == 2: 
+        if H > W:
+            pad_param = ((0, 0), (padd_1, padd_2))
+        else:
+            pad_param = ((padd_1, padd_2), (0, 0))
+    elif len(image.shape) == 3: 
+        if H > W:
+            pad_param = ((0, 0), (padd_1, padd_2), (0, 0))
+        else:
+            pad_param = ((padd_1, padd_2), (0, 0), (0, 0))
+
+    image = np.pad(image, pad_param, 'constant', constant_values=pad_value)
+
+    return image
+
+# Expand the image and mask
+def expand_image_mask(image, mask, ratio=1.4):
+    h,w = image.shape[0], image.shape[1]
+    H,W = int(h * ratio), int(w * ratio) 
+    h1 = int((H - h) // 2)
+    h2 = H - h - h1
+    w1 = int((W -w) // 2)
+    w2 = W -w - w1
+
+    pad_param_image = ((h1,h2),(w1,w2),(0,0))
+    pad_param_mask = ((h1,h2),(w1,w2))
+    image = np.pad(image, pad_param_image, 'constant', constant_values=255)
+    mask = np.pad(mask, pad_param_mask, 'constant', constant_values=0)
+    return image, mask
+
+# Convert the bounding box to a square shape
+def box2squre(image, box):
+    H,W = image.shape[0], image.shape[1]
+    y1,y2,x1,x2 = box
+    cx = (x1 + x2) // 2
+    cy = (y1 + y2) // 2
+    h,w = y2-y1, x2-x1
+
+    if h >= w:
+        x1 = cx - h//2
+        x2 = cx + h//2
+    else:
+        y1 = cy - w//2
+        y2 = cy + w//2
+    x1 = max(0,x1)
+    x2 = min(W,x2)
+    y1 = max(0,y1)
+    y2 = min(H,y2)
+    return (y1,y2,x1,x2)
+
+# Crop the predicted image back to the original image
+def crop_back( pred, tar_image,  extra_sizes, tar_box_yyxx_crop):
+    H1, W1, H2, W2 = extra_sizes
+    y1,y2,x1,x2 = tar_box_yyxx_crop    
+    pred = cv2.resize(pred, (W2, H2))
+    m = 2 # maigin_pixel
+
+    if W1 == H1:
+        if m != 0:
+            tar_image[y1+m :y2-m, x1+m:x2-m, :] =  pred[m:-m, m:-m]
+        else:
+            tar_image[y1 :y2, x1:x2, :] =  pred[:, :]
+        return tar_image
+
+    if W1 < W2:
+        pad1 = int((W2 - W1) / 2)
+        pad2 = W2 - W1 - pad1
+        pred = pred[:,pad1: -pad2, :]
+    else:
+        pad1 = int((H2 - H1) / 2)
+        pad2 = H2 - H1 - pad1
+        pred = pred[pad1: -pad2, :, :]
+
+    gen_image = tar_image.copy()
+    if m != 0:
+        gen_image[y1+m :y2-m, x1+m:x2-m, :] =  pred[m:-m, m:-m]
+    else:
+        gen_image[y1 :y2, x1:x2, :] =  pred[:, :]
+    
+    return gen_image
+