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)