first commit

.gitignore

@@ -0,0 +1,4 @@
+.vscode
+__pycache__
+experiments
+train_trm_coco.py

afftools/imutils.py

@@ -0,0 +1,280 @@
+
+import PIL.Image
+import random
+import numpy as np
+import pydensecrf.densecrf as dcrf
+from pydensecrf.utils import unary_from_softmax
+import os
+import cv2
+
+class RandomResizeLong():
+
+ def __init__(self, min_long, max_long):
+ self.min_long = min_long
+ self.max_long = max_long
+
+ def __call__(self, img, sal=None):
+
+ target_long = random.randint(self.min_long, self.max_long)
+ w, h = img.size
+
+ if w < h:
+ target_shape = (int(round(w * target_long / h)), target_long)
+ else:
+ target_shape = (target_long, int(round(h * target_long / w)))
+
+ img = img.resize(target_shape, resample=PIL.Image.CUBIC)
+ if sal:
+ sal = sal.resize(target_shape, resample=PIL.Image.CUBIC)
+ return img, sal
+ return img
+
+
+class RandomCrop():
+
+ def __init__(self, cropsize):
+ self.cropsize = cropsize
+
+ def __call__(self, imgarr, sal=None):
+
+ h, w, c = imgarr.shape
+
+ ch = min(self.cropsize, h)
+ cw = min(self.cropsize, w)
+
+ w_space = w - self.cropsize
+ h_space = h - self.cropsize
+
+ if w_space > 0:
+ cont_left = 0
+ img_left = random.randrange(w_space+1)
+ else:
+ cont_left = random.randrange(-w_space+1)
+ img_left = 0
+
+ if h_space > 0:
+ cont_top = 0
+ img_top = random.randrange(h_space+1)
+ else:
+ cont_top = random.randrange(-h_space+1)
+ img_top = 0
+
+ container = np.zeros((self.cropsize, self.cropsize, imgarr.shape[-1]), np.float32)
+ container[cont_top:cont_top+ch, cont_left:cont_left+cw] = \
+ imgarr[img_top:img_top+ch, img_left:img_left+cw]
+ if sal is not None:
+ container_sal = np.zeros((self.cropsize, self.cropsize,1), np.float32)
+ container_sal[cont_top:cont_top+ch, cont_left:cont_left+cw,0] = \
+ sal[img_top:img_top+ch, img_left:img_left+cw]
+ return container, container_sal
+
+ return container
+
+def get_random_crop_box(imgsize, cropsize):
+ h, w = imgsize
+
+ ch = min(cropsize, h)
+ cw = min(cropsize, w)
+
+ w_space = w - cropsize
+ h_space = h - cropsize
+
+ if w_space > 0:
+ cont_left = 0
+ img_left = random.randrange(w_space + 1)
+ else:
+ cont_left = random.randrange(-w_space + 1)
+ img_left = 0
+
+ if h_space > 0:
+ cont_top = 0
+ img_top = random.randrange(h_space + 1)
+ else:
+ cont_top = random.randrange(-h_space + 1)
+ img_top = 0
+
+ return cont_top, cont_top+ch, cont_left, cont_left+cw, img_top, img_top+ch, img_left, img_left+cw
+
+def crop_with_box(img, box):
+ if len(img.shape) == 3:
+ img_cont = np.zeros((max(box[1]-box[0], box[4]-box[5]), max(box[3]-box[2], box[7]-box[6]), img.shape[-1]), dtype=img.dtype)
+ else:
+ img_cont = np.zeros((max(box[1] - box[0], box[4] - box[5]), max(box[3] - box[2], box[7] - box[6])), dtype=img.dtype)
+ img_cont[box[0]:box[1], box[2]:box[3]] = img[box[4]:box[5], box[6]:box[7]]
+ return img_cont
+
+
+def random_crop(images, cropsize, fills):
+ if isinstance(images[0], PIL.Image.Image):
+ imgsize = images[0].size[::-1]
+ else:
+ imgsize = images[0].shape[:2]
+ box = get_random_crop_box(imgsize, cropsize)
+
+ new_images = []
+ for img, f in zip(images, fills):
+
+ if isinstance(img, PIL.Image.Image):
+ img = img.crop((box[6], box[4], box[7], box[5]))
+ cont = PIL.Image.new(img.mode, (cropsize, cropsize))
+ cont.paste(img, (box[2], box[0]))
+ new_images.append(cont)
+
+ else:
+ if len(img.shape) == 3:
+ cont = np.ones((cropsize, cropsize, img.shape[2]), img.dtype)*f
+ else:
+ cont = np.ones((cropsize, cropsize), img.dtype)*f
+ cont[box[0]:box[1], box[2]:box[3]] = img[box[4]:box[5], box[6]:box[7]]
+ new_images.append(cont)
+
+ return new_images
+
+
+class AvgPool2d():
+
+ def __init__(self, ksize):
+ self.ksize = ksize
+
+ def __call__(self, img):
+ import skimage.measure
+
+ return skimage.measure.block_reduce(img, (self.ksize, self.ksize, 1), np.mean)
+
+
+class RandomHorizontalFlip():
+ def __init__(self):
+ return
+
+ def __call__(self, img, sal=None):
+ if bool(random.getrandbits(1)):
+ #img = img.transpose(PIL.Image.FLIP_LEFT_RIGHT)
+ img = np.fliplr(img).copy()
+ if sal:
+ #sal = sal.transpose(PIL.Image.FLIP_LEFT_RIGHT)
+ sal = np.fliplr(sal).copy()
+ return img, sal 
+ return img
+ else:
+ if sal:
+ return img, sal
+ return img
+
+
+class CenterCrop():
+
+ def __init__(self, cropsize, default_value=0):
+ self.cropsize = cropsize
+ self.default_value = default_value
+
+ def __call__(self, npimg):
+
+ h, w = npimg.shape[:2]
+
+ ch = min(self.cropsize, h)
+ cw = min(self.cropsize, w)
+
+ sh = h - self.cropsize
+ sw = w - self.cropsize
+
+ if sw > 0:
+ cont_left = 0
+ img_left = int(round(sw / 2))
+ else:
+ cont_left = int(round(-sw / 2))
+ img_left = 0
+
+ if sh > 0:
+ cont_top = 0
+ img_top = int(round(sh / 2))
+ else:
+ cont_top = int(round(-sh / 2))
+ img_top = 0
+
+ if len(npimg.shape) == 2:
+ container = np.ones((self.cropsize, self.cropsize), npimg.dtype)*self.default_value
+ else:
+ container = np.ones((self.cropsize, self.cropsize, npimg.shape[2]), npimg.dtype)*self.default_value
+
+ container[cont_top:cont_top+ch, cont_left:cont_left+cw] = \
+ npimg[img_top:img_top+ch, img_left:img_left+cw]
+
+ return container
+
+
+def HWC_to_CHW(tensor, sal=False):
+ if sal:
+ tensor = np.expand_dims(tensor, axis=0)
+ else:
+ tensor = np.transpose(tensor, (2, 0, 1))
+ return tensor
+
+
+class RescaleNearest():
+ def __init__(self, scale):
+ self.scale = scale
+
+ def __call__(self, npimg):
+ import cv2
+ return cv2.resize(npimg, None, fx=self.scale, fy=self.scale, interpolation=cv2.INTER_NEAREST)
+
+
+
+
+def crf_inference(img, probs, t=10, scale_factor=1, labels=21):
+ import pydensecrf.densecrf as dcrf
+ from pydensecrf.utils import unary_from_softmax
+
+ h, w = img.shape[:2]
+ n_labels = labels
+
+ d = dcrf.DenseCRF2D(w, h, n_labels)
+
+ unary = unary_from_softmax(probs)
+ unary = np.ascontiguousarray(unary)
+
+ d.setUnaryEnergy(unary)
+ d.addPairwiseGaussian(sxy=3/scale_factor, compat=3)
+ d.addPairwiseBilateral(sxy=80/scale_factor, srgb=13, rgbim=np.copy(img), compat=10)
+ Q = d.inference(t)
+
+ return np.array(Q).reshape((n_labels, h, w))
+
+def crf_inference_ysh(name, probs, t=10, scale_factor=1, labels=21):
+
+ image_path = os.path.join('./data/VOC2012/JPEGImages', name + '.jpg')
+ img = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(np.float32)
+ img = np.ascontiguousarray(np.uint8(img))
+
+ h, w = img.shape[:2]
+ n_labels = labels
+
+ d = dcrf.DenseCRF2D(w, h, n_labels)
+
+ unary = unary_from_softmax(probs)
+ unary = np.ascontiguousarray(unary)
+
+ d.setUnaryEnergy(unary)
+
+ # d.addPairwiseGaussian(sxy=1, compat=3) # 64.85
+ # d.addPairwiseBilateral(sxy=67, srgb=3, rgbim=np.copy(img), compat=4) # 64.85
+
+ d.addPairwiseGaussian(sxy=1, compat=3) # 64.86/no normal
+ d.addPairwiseBilateral(sxy=67, srgb=3, rgbim=np.copy(img), compat=4) #64.86/no normal
+
+ # d.addPairwiseGaussian(sxy=3, compat=3) #RRM
+ # d.addPairwiseBilateral(sxy=80, srgb=13, rgbim=np.copy(img), compat=10) #RRM t=10
+
+ # d.addPairwiseGaussian(sxy=1, compat=3) #good
+ # d.addPairwiseBilateral(sxy=67, srgb=3, rgbim=np.copy(img), compat=4) #good
+
+ # d.addPairwiseGaussian(sxy=2, compat=3,kernel=dcrf.DIAG_KERNEL,normalization=dcrf.NORMALIZE_SYMMETRIC)
+ # d.addPairwiseBilateral(sxy=80, srgb=13, rgbim=np.copy(img),kernel=dcrf.DIAG_KERNEL, normalization=dcrf.NORMALIZE_SYMMETRIC,compat=4) #good
+
+ # d.addPairwiseGaussian(sxy=2, compat=3, kernel=dcrf.DIAG_KERNEL, normalization=dcrf.NORMALIZE_SYMMETRIC)
+ # d.addPairwiseBilateral(sxy=80, srgb=13, rgbim=np.copy(img), kernel=dcrf.DIAG_KERNEL,normalization=dcrf.NORMALIZE_SYMMETRIC, compat=4) # good
+
+
+ Q = d.inference(t)
+
+ return np.array(Q).reshape((n_labels, h, w))