def loaddata():
video_index = random.randint(0, len(videos) - 1)
video = videos[video_index]
img_index = random.randint(N, lengths[video_index] - N)
input_img = np.zeros((SIZE, SIZE, 3 * N + 1), dtype='uint8')
for i in range(0, N):
# print('./dataset/'+video+'/mosaic/output_'+'%05d'%(img_index+i-int(N/2))+'.png')
img = cv2.imread('./dataset/' + video + '/mosaic/output_' + '%05d' %
(img_index + i - int(N / 2)) + '.png')
img = impro.resize(img, SIZE)
input_img[:, :, i * 3:(i + 1) * 3] = img
mask = cv2.imread(
'./dataset/' + video + '/mask/output_' + '%05d' % (img_index) + '.png',
0)
mask = impro.resize(mask, 256)
mask = impro.mask_threshold(mask, 15, 128)
input_img[:, :, -1] = mask
input_img = data.im2tensor(input_img,
bgr2rgb=False,
use_gpu=opt.use_gpu,
use_transform=False)
ground_true = cv2.imread('./dataset/' + video + '/ori/output_' + '%05d' %
(img_index) + '.png')
ground_true = impro.resize(ground_true, SIZE)
# ground_true = im2tensor(ground_true,use_gpu)
ground_true = data.im2tensor(ground_true,
bgr2rgb=False,
use_gpu=opt.use_gpu,
use_transform=False)
return input_img, ground_true
def get_mosaic_position(img_origin, net_mosaic_pos, opt):
mask = run_unet_rectim(img_origin, net_mosaic_pos, use_gpu=opt.use_gpu)
mask = impro.mask_threshold(mask, 10, 128)
x, y, size, area = impro.boundingSquare(mask, Ex_mul=1.5)
rat = min(img_origin.shape[:2]) / 128.0
x, y, size = int(rat * x), int(rat * y), int(rat * size)
return x, y, size
def get_ROI_position(img, net, opt, keepsize=True):
mask = run_segment(img, net, size=360, gpu_id=opt.gpu_id)
mask = impro.mask_threshold(mask, opt.mask_extend, opt.mask_threshold)
if keepsize:
mask = impro.resize_like(mask, img)
x, y, halfsize, area = impro.boundingSquare(mask, 1)
return mask, x, y, halfsize, area
def get_mosaic_position(img_origin,net_mosaic_pos,opt,threshold = 128 ):
mask = run_unet(img_origin,net_mosaic_pos,size=224,use_gpu = opt.use_gpu)
mask = impro.mask_threshold(mask,30,threshold)
if not opt.all_mosaic_area:
mask = impro.find_mostlikely_ROI(mask)
x,y,size,area = impro.boundingSquare(mask,Ex_mul=opt.ex_mult)
rat = min(img_origin.shape[:2])/224.0
x,y,size = int(rat*x),int(rat*y),int(rat*size)
return x,y,size,mask
def get_mosaic_position(img_origin, net_mosaic_pos, opt, threshold=128):
mask = run_unet_rectim(img_origin, net_mosaic_pos, use_gpu=opt.use_gpu)
#mask_1 = mask.copy()
mask = impro.mask_threshold(mask, 30, threshold)
if not opt.no_large_area:
mask = impro.find_best_ROI(mask)
x, y, size, area = impro.boundingSquare(mask, Ex_mul=opt.ex_mult)
rat = min(img_origin.shape[:2]) / 224.0
x, y, size = int(rat * x), int(rat * y), int(rat * size)
return x, y, size, mask
def get_mosaic_position(img_origin, net_mosaic_pos, opt):
h, w = img_origin.shape[:2]
mask = run_segment(img_origin, net_mosaic_pos, size=360, gpu_id=opt.gpu_id)
# mask_1 = mask.copy()
mask = impro.mask_threshold(mask,
ex_mun=int(min(h, w) / 20),
threshold=opt.mask_threshold)
if not opt.all_mosaic_area:
mask = impro.find_mostlikely_ROI(mask)
x, y, size, area = impro.boundingSquare(mask, Ex_mul=opt.ex_mult)
#Location fix
rat = min(h, w) / 360.0
x, y, size = int(rat * x), int(rat * y), int(rat * size)
x, y = np.clip(x, 0, w), np.clip(y, 0, h)
size = np.clip(size, 0, min(w - x, h - y))
# print(x,y,size)
return x, y, size, mask
def get_ROI_position(img, net, opt):
mask = run_unet_rectim(img, net, use_gpu=opt.use_gpu)
mask = impro.mask_threshold(mask, opt.mask_extend, opt.mask_threshold)
x, y, halfsize, area = impro.boundingSquare(mask, 1)
return mask, x, y, area
#test
if os.path.isdir('./test'):
netG.eval()
test_names = os.listdir('./test')
test_names.sort()
result = np.zeros((opt.finesize*2,opt.finesize*len(test_names),3), dtype='uint8')
for cnt,test_name in enumerate(test_names,0):
img_names = os.listdir(os.path.join('./test',test_name,'image'))
img_names.sort()
inputdata = np.zeros((opt.finesize,opt.finesize,3*N+1), dtype='uint8')
for i in range(0,N):
img = impro.imread(os.path.join('./test',test_name,'image',img_names[i]))
img = impro.resize(img,opt.finesize)
inputdata[:,:,i*3:(i+1)*3] = img
mask = impro.imread(os.path.join('./test',test_name,'mask.png'),'gray')
mask = impro.resize(mask,opt.finesize)
mask = impro.mask_threshold(mask,15,128)
inputdata[:,:,-1] = mask
result[0:opt.finesize,opt.finesize*cnt:opt.finesize*(cnt+1),:] = inputdata[:,:,int((N-1)/2)*3:(int((N-1)/2)+1)*3]
inputdata = data.im2tensor(inputdata,bgr2rgb=False,use_gpu=opt.use_gpu,use_transform = False,is0_1 = False)
pred = netG(inputdata)
pred = data.tensor2im(pred,rgb2bgr = False, is0_1 = False)
result[opt.finesize:opt.finesize*2,opt.finesize*cnt:opt.finesize*(cnt+1),:] = pred
cv2.imwrite(os.path.join(dir_checkpoint,str(iter+1)+'_test.jpg'), result)
netG.train()
for imagepath in imagepaths:
img = impro.imread(imagepath)
mask = runmodel.get_ROI_position(img, net, opt,
keepsize=True)[0]
imgs.append(img)
masks.append(mask)
if not mask_flag:
mask_avg = mask.astype(np.float64)
mask_flag = True
else:
mask_avg += mask.astype(np.float64)
mask_avg = np.clip(mask_avg / len(imagepaths), 0,
255).astype('uint8')
mask_avg = impro.mask_threshold(mask_avg, 20, 64)
if not opt.all_mosaic_area:
mask_avg = impro.find_mostlikely_ROI(mask_avg)
x, y, size, area = impro.boundingSquare(mask_avg,
Ex_mul=random.uniform(
1.1, 1.5))
for i in range(len(imagepaths)):
img = impro.resize(imgs[i][y - size:y + size,
x - size:x + size],
opt.outsize,
interpolation=cv2.INTER_CUBIC)
mask = impro.resize(masks[i][y - size:y + size,
x - size:x + size],
opt.outsize,
interpolation=cv2.INTER_CUBIC)
def add_mosaic_to_image(path):
img = cv2.imread(path)
mask =runmodel.run_unet_rectim(img,net,use_gpu = opt.use_gpu)
mask = impro.mask_threshold(mask,opt.mask_extend,opt.mask_threshold)
img = mosaic.addmosaic(img,mask,opt.mosaic_size,opt.output_size,model = opt.mosaic_mod)
return img
imagepath = os.path.join('./tmp/video2image', imagepath)
print('Find ROI location:', imagepath)
img = impro.imread(imagepath)
x, y, size, mask = runmodel.get_mosaic_position(img,
net,
opt,
threshold=64)
cv2.imwrite(
os.path.join('./tmp/ROI_mask', os.path.basename(imagepath)), mask)
positions.append([x, y, size])
mask_avg = mask_avg + mask
print('Optimize ROI locations...')
mask_index = filt.position_medfilt(np.array(positions), 13)
mask = np.clip(mask_avg / len(imagepaths), 0, 255).astype('uint8')
mask = impro.mask_threshold(mask, 20, 32)
x, y, size, area = impro.boundingSquare(mask, Ex_mul=1.5)
rat = min(img_ori_example.shape[:2]) / 128.0
x, y, size = int(rat * x), int(rat * y), int(rat * size)
cv2.imwrite(os.path.join('./tmp/ROI_mask_check', 'test_show.png'), mask)
if size != 0:
mask_path = './dataset/' + os.path.splitext(
os.path.basename(path))[0] + '/mask'
ori_path = './dataset/' + os.path.splitext(
os.path.basename(path))[0] + '/ori'
mosaic_path = './dataset/' + os.path.splitext(
os.path.basename(path))[0] + '/mosaic'
os.makedirs('./dataset/' +
os.path.splitext(os.path.basename(path))[0] + '')
os.makedirs(mask_path)
os.makedirs(ori_path)
