torchbenchmark/models/Background_Matting/data_loader.py

from __future__ import print_function, division
import os
import torch
import pandas as pd
import skimage
from skimage import io
import numpy as np
import matplotlib.pyplot as plt
import pdb, random
from torch.utils.data import Dataset, DataLoader
import random, os, cv2

unknown_code=128

class VideoData(Dataset):
	def __init__(self,csv_file,data_config,transform=None):
		self.frames = pd.read_csv(csv_file,sep=';')
		self.transform = transform
		self.resolution=data_config['reso']
		
	def __len__(self):
		return len(self.frames)

	def __getitem__(self,idx):
		img = io.imread(self.frames.iloc[idx, 0])
		back = io.imread(self.frames.iloc[idx, 1])
		seg = io.imread(self.frames.iloc[idx, 2])
		
		fr1 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 3]), cv2.COLOR_BGR2GRAY)
		fr2 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 4]), cv2.COLOR_BGR2GRAY)
		fr3 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 5]), cv2.COLOR_BGR2GRAY)
		fr4 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 6]), cv2.COLOR_BGR2GRAY)

		back_rnd = io.imread(self.frames.iloc[idx, 7])
		
		sz=self.resolution

		if np.random.random_sample() > 0.5:
			img = cv2.flip(img,1)
			seg = cv2.flip(seg,1)
			back = cv2.flip(back,1)
			back_rnd = cv2.flip(back_rnd,1)
			fr1=cv2.flip(fr1,1); fr2=cv2.flip(fr2,1); fr3=cv2.flip(fr3,1); fr4=cv2.flip(fr4,1)

		#make frames together
		multi_fr=np.zeros((img.shape[0],img.shape[1],4))
		multi_fr[...,0]=fr1; multi_fr[...,1]=fr2; multi_fr[...,2]=fr3; multi_fr[...,3]=fr4;
			
		
		#allow random cropping centered on the segmentation map
		bbox=create_bbox(seg,seg.shape[0],seg.shape[1])
		img=apply_crop(img,bbox,self.resolution)
		seg=apply_crop(seg,bbox,self.resolution)
		back=apply_crop(back,bbox,self.resolution)
		back_rnd=apply_crop(back_rnd,bbox,self.resolution)
		multi_fr=apply_crop(multi_fr,bbox,self.resolution)

		#convert seg to guidance map
		#segg=create_seg_guide(seg,self.resolution)

		sample = {'image': to_tensor(img), 'seg': to_tensor(create_seg_guide(seg,self.resolution)), 'bg': to_tensor(back), 'multi_fr': to_tensor(multi_fr), 'seg-gt':to_tensor(seg), 'back-rnd': to_tensor(back_rnd)}

		if self.transform:
			sample = self.transform(sample)
		return sample


class AdobeDataAffineHR(Dataset):
	def __init__(self,csv_file,data_config,transform=None):
		self.frames = pd.read_csv(csv_file,sep=';')
		self.transform = transform
		self.resolution=data_config['reso']
		self.trimapK=data_config['trimapK']
		self.noise=data_config['noise']
		
	def __len__(self):
		return len(self.frames)

	def __getitem__(self,idx):
		try:
			#load
			fg = io.imread(self.frames.iloc[idx, 0])
			alpha = io.imread(self.frames.iloc[idx, 1])
			image = io.imread(self.frames.iloc[idx, 2])
			back = io.imread(self.frames.iloc[idx, 3])

			fg = cv2.resize(fg, dsize=(800,800))
			alpha = cv2.resize(alpha, dsize=(800,800))
			back = cv2.resize(back, dsize=(800,800))
			image = cv2.resize(image, dsize=(800,800))


			sz=self.resolution

			#random flip
			if np.random.random_sample() > 0.5:
				alpha = cv2.flip(alpha,1)
				fg = cv2.flip(fg,1)
				back = cv2.flip(back,1)
				image = cv2.flip(image,1)

			trimap=generate_trimap(alpha,self.trimapK[0],self.trimapK[1],False)


			#randcom crop+scale
			different_sizes = [(576,576),(608,608),(640,640),(672,672),(704,704),(736,736),(768,768),(800,800)]
			crop_size = random.choice(different_sizes)

			x, y = random_choice(trimap, crop_size)

			fg = safe_crop(fg, x, y, crop_size,sz)
			alpha = safe_crop(alpha, x, y, crop_size,sz)
			image = safe_crop(image, x, y, crop_size,sz)
			back = safe_crop(back, x, y, crop_size,sz)
			trimap = safe_crop(trimap, x, y, crop_size,sz)

			#Perturb Background: random noise addition or gamma change
			if self.noise:
				if np.random.random_sample() > 0.6:
					sigma=np.random.randint(low=2, high=6)
					mu=np.random.randint(low=0, high=14)-7
					back_tr=add_noise(back,mu,sigma)
				else:
					back_tr=skimage.exposure.adjust_gamma(back,np.random.normal(1,0.12))


			#Create motion cues: transform foreground and create 4 additional images
			affine_fr=np.zeros((fg.shape[0],fg.shape[1],4))
			for t in range(0,4):
				T=np.random.normal(0,5,(2,1)); theta=np.random.normal(0,7);
				R=np.array([[np.cos(np.deg2rad(theta)), -np.sin(np.deg2rad(theta))],[np.sin(np.deg2rad(theta)), np.cos(np.deg2rad(theta))]])
				sc=np.array([[1+np.random.normal(0,0.05), 0],[0,1]]); sh=np.array([[1, np.random.normal(0,0.05)*(np.random.random_sample() > 0.5)],[np.random.normal(0,0.05)*(np.random.random_sample() > 0.5), 1]]);
				A=np.concatenate((sc*sh*R, T), axis=1);

				fg_tr = cv2.warpAffine(fg.astype(np.uint8),A,(fg.shape[1],fg.shape[0]),flags=cv2.INTER_LINEAR,borderMode=cv2.BORDER_REFLECT)
				alpha_tr = cv2.warpAffine(alpha.astype(np.uint8),A,(fg.shape[1],fg.shape[0]),flags=cv2.INTER_NEAREST,borderMode=cv2.BORDER_REFLECT)

				sigma=np.random.randint(low=2, high=6)
				mu=np.random.randint(low=0, high=14)-7
				back_tr0=add_noise(back,mu,sigma)

				affine_fr[...,t]=cv2.cvtColor(composite(fg_tr,back_tr0,alpha_tr), cv2.COLOR_BGR2GRAY)



			sample = {'image': to_tensor(image), 'fg': to_tensor(fg), 'alpha': to_tensor(alpha), 'bg': to_tensor(back), 'trimap': to_tensor(trimap), 'bg_tr': to_tensor(back_tr), 'seg': to_tensor(create_seg(alpha,trimap)), 'multi_fr': to_tensor(affine_fr)}



			if self.transform:
				sample = self.transform(sample)
			return sample
		except Exception as e:
			print("Error loading: " + self.frames.iloc[idx, 0])
			print(e)


#Functions

def create_seg_guide(rcnn,reso):
	kernel_er = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
	kernel_dil = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
	rcnn=rcnn.astype(np.float32)/255; rcnn[rcnn>0.2]=1;
	K=25

	zero_id=np.nonzero(np.sum(rcnn,axis=1)==0)
	del_id=zero_id[0][zero_id[0]>250]
	if len(del_id)>0:
		del_id=[del_id[0]-2,del_id[0]-1,*del_id]
		rcnn=np.delete(rcnn,del_id,0)
	rcnn = cv2.copyMakeBorder( rcnn, 0, K + len(del_id), 0, 0, cv2.BORDER_REPLICATE)

	rcnn = cv2.erode(rcnn, kernel_er, iterations=np.random.randint(10,20))
	rcnn = cv2.dilate(rcnn, kernel_dil, iterations=np.random.randint(3,7))
	k_size_list=[(21,21),(31,31),(41,41)]
	rcnn=cv2.GaussianBlur(rcnn.astype(np.float32),random.choice(k_size_list),0)
	rcnn=(255*rcnn).astype(np.uint8)
	rcnn=np.delete(rcnn, range(reso[0],reso[0]+K), 0)

	return rcnn

def crop_holes(img,cx,cy,crop_size):
	img[cy:cy+crop_size[0],cx:cx+crop_size[1]]=0
	return img

def create_seg(alpha,trimap):
	#old
	num_holes=np.random.randint(low=0, high=3)
	crop_size_list=[(15,15),(25,25),(35,35),(45,45)]
	kernel_er = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
	kernel_dil = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
	seg = (alpha>0.5).astype(np.float32)
	#print('Before %.4f max: %.4f' %(seg.sum(),seg.max()))
	#old
	seg = cv2.erode(seg, kernel_er, iterations=np.random.randint(low=10,high=20))
	seg = cv2.dilate(seg, kernel_dil, iterations=np.random.randint(low=15,high=30))
	#print('After %.4f max: %.4f' %(seg.sum(),seg.max()))
	seg=seg.astype(np.float32)
	seg=(255*seg).astype(np.uint8)
	for i in range(num_holes):
		crop_size=random.choice(crop_size_list)
		cx,cy = random_choice(trimap,crop_size)
		seg=crop_holes(seg,cx,cy,crop_size)
		trimap=crop_holes(trimap,cx,cy,crop_size)
	k_size_list=[(21,21),(31,31),(41,41)]
	seg=cv2.GaussianBlur(seg.astype(np.float32),random.choice(k_size_list),0)
	return seg.astype(np.uint8)


def apply_crop(img,bbox,reso):
	img_crop=img[bbox[0]:bbox[0]+bbox[2],bbox[1]:bbox[1]+bbox[3],...]; 
	img_crop=cv2.resize(img_crop,reso)
	return img_crop

def create_bbox(mask,R,C):
	where = np.array(np.where(mask))
	x1, y1 = np.amin(where, axis=1)
	x2, y2 = np.amax(where, axis=1)

	w=np.maximum(y2-y1,x2-x1);
	bd=np.random.uniform(0.1,0.4)
	x1=x1-np.round(bd*w)
	y1=y1-np.round(bd*w)
	y2=y2+np.round(bd*w)
	
	if x1<0: x1=0
	if y1<0: y1=0
	if y2>=C: y2=C
	if x2>=R: x2=R-1
	
	bbox=np.around([x1,y1,x2-x1,y2-y1]).astype('int')

	return bbox

def composite(fg, bg, a):
	fg = fg.astype(np.float32); bg=bg.astype(np.float32); a=a.astype(np.float32);
	alpha= np.expand_dims(a / 255,axis=2)
	im = alpha * fg + (1 - alpha) * bg
	im = im.astype(np.uint8)
	return im

def add_noise(back,mean,sigma):
	back=back.astype(np.float32)
	row,col,ch= back.shape
	gauss = np.random.normal(mean,sigma,(row,col,ch))
	gauss = gauss.reshape(row,col,ch)
	#gauss = np.repeat(gauss[:, :, np.newaxis], ch, axis=2)
	noisy = back + gauss

	noisy[noisy<0]=0; noisy[noisy>255]=255;

	return noisy.astype(np.uint8)

def safe_crop(mat, x, y, crop_size,img_size,cubic=True):
	img_rows, img_cols = img_size
	crop_height, crop_width = crop_size
	if len(mat.shape) == 2:
		ret = np.zeros((crop_height, crop_width), np.float32)
	else:
		ret = np.zeros((crop_height, crop_width, 3), np.float32)
	crop = mat[y:y + crop_height, x:x + crop_width]
	h, w = crop.shape[:2]
	ret[0:h, 0:w] = crop
	if crop_size != (img_rows, img_cols):
		if cubic:
			ret = cv2.resize(ret, dsize=(img_rows, img_cols))
		else:
			ret = cv2.resize(ret, dsize=(img_rows, img_cols), interpolation=cv2.INTER_NEAREST)
	return ret

def generate_trimap(alpha,K1,K2,train_mode):
	kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
	fg = np.array(np.equal(alpha, 255).astype(np.float32))
	if train_mode:
		K=np.random.randint(K1,K2)
	else:
		K=np.round((K1+K2)/2).astype('int')

	fg = cv2.erode(fg, kernel, iterations=K)
	unknown = np.array(np.not_equal(alpha, 0).astype(np.float32))
	unknown = cv2.dilate(unknown, kernel, iterations=2*K)
	trimap = fg * 255 + (unknown - fg) * 128
	return trimap.astype(np.uint8)


def random_choice(trimap, crop_size=(320, 320)):
	img_height, img_width = trimap.shape[0:2]
	crop_height, crop_width = crop_size

	val_idx=np.zeros((img_height,img_width))
	val_idx[int(crop_height/2):int(img_height-crop_height/2),int(crop_width/2):int(img_width-crop_width/2)]=1

	y_indices, x_indices = np.where(np.logical_and(trimap == unknown_code,val_idx==1))
	num_unknowns = len(y_indices)
	x, y = 0, 0
	if num_unknowns > 0:
		ix = np.random.choice(range(num_unknowns))
		center_x = x_indices[ix]
		center_y = y_indices[ix]
		x = max(0, center_x - int(crop_width / 2))
		y = max(0, center_y - int(crop_height / 2))

	#added extra

	return x, y

def to_tensor(pic):
	if len(pic.shape)>=3:
		img = torch.from_numpy(pic.transpose((2, 0, 1)))
	else:
		img=torch.from_numpy(pic)
		img=img.unsqueeze(0)
	# backward compatibility


	return 2*(img.float().div(255))-1