File size: 2,173 Bytes
352cafd |
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 |
import cv2
import numpy as np
def get_random_structure(size):
# The provided model is trained with
# choice = np.random.randint(4)
# instead, which is a bug that we fixed here
choice = np.random.randint(1, 5)
if choice == 1:
return cv2.getStructuringElement(cv2.MORPH_RECT, (size, size))
elif choice == 2:
return cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size))
elif choice == 3:
return cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size//2))
elif choice == 4:
return cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size//2, size))
def random_dilate(seg, min=3, max=10):
size = np.random.randint(min, max)
kernel = get_random_structure(size)
seg = cv2.dilate(seg,kernel,iterations = 1)
return seg
def random_erode(seg, min=3, max=10):
size = np.random.randint(min, max)
kernel = get_random_structure(size)
seg = cv2.erode(seg,kernel,iterations = 1)
return seg
def compute_iou(seg, gt):
intersection = seg*gt
union = seg+gt
return (np.count_nonzero(intersection) + 1e-6) / (np.count_nonzero(union) + 1e-6)
def perturb_seg(gt, iou_target=0.6):
h, w = gt.shape
seg = gt.copy()
_, seg = cv2.threshold(seg, 127, 255, 0)
# Rare case
if h <= 2 or w <= 2:
print('GT too small, returning original')
return seg
# Do a bunch of random operations
for _ in range(250):
for _ in range(4):
lx, ly = np.random.randint(w), np.random.randint(h)
lw, lh = np.random.randint(lx+1,w+1), np.random.randint(ly+1,h+1)
# Randomly set one pixel to 1/0. With the following dilate/erode, we can create holes/external regions
if np.random.rand() < 0.25:
cx = int((lx + lw) / 2)
cy = int((ly + lh) / 2)
seg[cy, cx] = np.random.randint(2) * 255
if np.random.rand() < 0.5:
seg[ly:lh, lx:lw] = random_dilate(seg[ly:lh, lx:lw])
else:
seg[ly:lh, lx:lw] = random_erode(seg[ly:lh, lx:lw])
if compute_iou(seg, gt) < iou_target:
break
return seg
|