def cleanmosaic_video_byframe(opt,netG,netM):
path = opt.media_path
fps,imagepaths = video_init(opt,path)
positions = []
# get position
for i,imagepath in enumerate(imagepaths,1):
img_origin = impro.imread(os.path.join('./tmp/video2image',imagepath))
x,y,size = runmodel.get_mosaic_position(img_origin,netM,opt)[:3]
positions.append([x,y,size])
print('\r','Find mosaic location:'+str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='')
print('\nOptimize mosaic locations...')
positions =np.array(positions)
for i in range(3):positions[:,i] = filt.medfilt(positions[:,i],opt.medfilt_num)
# clean mosaic
for i,imagepath in enumerate(imagepaths,0):
x,y,size = positions[i][0],positions[i][1],positions[i][2]
img_origin = impro.imread(os.path.join('./tmp/video2image',imagepath))
img_result = img_origin.copy()
if size != 0:
img_mosaic = img_origin[y-size:y+size,x-size:x+size]
img_fake = runmodel.run_pix2pix(img_mosaic,netG,opt)
img_result = impro.replace_mosaic(img_origin,img_fake,x,y,size,opt.no_feather)
cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_result)
print('\r','Clean Mosaic:'+str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='')
print()
ffmpeg.image2video( fps,
'./tmp/replace_mosaic/output_%05d.'+opt.tempimage_type,
'./tmp/voice_tmp.mp3',
os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_clean.mp4'))
def loadimage(dir_img, dir_mask, loadsize, eval_p):
t1 = datetime.datetime.now()
imgnames = os.listdir(dir_img)
# imgnames = imgnames[:100]
random.shuffle(imgnames)
imgnames = imgnames[:MAX_LOAD]
print('load images:', len(imgnames))
imgnames = (f[:-4] for f in imgnames)
images = []
masks = []
for imgname in imgnames:
img = impro.imread(dir_img + imgname + '.jpg')
mask = impro.imread(dir_mask + imgname + '.png', mod='gray')
img = impro.resize(img, loadsize)
mask = impro.resize(mask, loadsize)
images.append(img)
masks.append(mask)
train_images, train_masks = images[0:int(len(masks) * (
1 - eval_p))], masks[0:int(len(masks) * (1 - eval_p))]
eval_images, eval_masks = images[int(len(masks) * (
1 - eval_p)):len(masks)], masks[int(len(masks) *
(1 - eval_p)):len(masks)]
t2 = datetime.datetime.now()
print('load data cost time:', (t2 - t1).seconds, 's')
return train_images, train_masks, eval_images, eval_masks
def addmosaic_video(opt,netS):
path = opt.media_path
fps,imagepaths = video_init(opt,path)
# get position
positions = []
for i,imagepath in enumerate(imagepaths,1):
img = impro.imread(os.path.join('./tmp/video2image',imagepath))
mask,x,y,area = runmodel.get_ROI_position(img,netS,opt)
positions.append([x,y,area])
cv2.imwrite(os.path.join('./tmp/ROI_mask',imagepath),mask)
print('\r','Find ROI location:'+str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='')
print('\nOptimize ROI locations...')
mask_index = filt.position_medfilt(np.array(positions), 7)
# add mosaic
for i in range(len(imagepaths)):
mask = impro.imread(os.path.join('./tmp/ROI_mask',imagepaths[mask_index[i]]),'gray')
img = impro.imread(os.path.join('./tmp/video2image',imagepaths[i]))
if impro.mask_area(mask)>100:
img = mosaic.addmosaic(img, mask, opt)
cv2.imwrite(os.path.join('./tmp/addmosaic_image',imagepaths[i]),img)
print('\r','Add Mosaic:'+str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='')
print()
ffmpeg.image2video( fps,
'./tmp/addmosaic_image/output_%05d.'+opt.tempimage_type,
'./tmp/voice_tmp.mp3',
os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_add.mp4'))
def addmosaic_video(opt, netS):
path = opt.media_path
fps, imagepaths = video_init(opt, path)[:2]
length = len(imagepaths)
start_frame = int(imagepaths[0][7:13])
mask_index = get_roi_positions(opt, netS, imagepaths)[(start_frame - 1):]
t1 = time.time()
if not opt.no_preview:
cv2.namedWindow('preview', cv2.WINDOW_NORMAL)
# add mosaic
print('Step:3/4 -- Add Mosaic:')
t1 = time.time()
# print(mask_index)
for i, imagepath in enumerate(imagepaths, 1):
mask = impro.imread(
os.path.join(
opt.temp_dir + '/ROI_mask',
imagepaths[np.clip(mask_index[i - 1] - start_frame, 0,
1000000)]), 'gray')
img = impro.imread(
os.path.join(opt.temp_dir + '/video2image', imagepath))
if impro.mask_area(mask) > 100:
try: #Avoid unknown errors
img = mosaic.addmosaic(img, mask, opt)
except Exception as e:
print('Warning:', e)
t = Thread(target=cv2.imwrite,
args=(os.path.join(opt.temp_dir + '/addmosaic_image',
imagepath), img))
t.start()
os.remove(os.path.join(opt.temp_dir + '/video2image', imagepath))
#preview result and print
if not opt.no_preview:
cv2.imshow('preview', img)
cv2.waitKey(1) & 0xFF
t2 = time.time()
print('\r',
str(i) + '/' + str(length),
util.get_bar(100 * i / length, num=35),
util.counttime(t1, t2, i, length),
end='')
print()
if not opt.no_preview:
cv2.destroyAllWindows()
print('Step:4/4 -- Convert images to video')
ffmpeg.image2video(
fps,
opt.temp_dir + '/addmosaic_image/output_%06d.' + opt.tempimage_type,
opt.temp_dir + '/voice_tmp.mp3',
os.path.join(opt.result_dir,
os.path.splitext(os.path.basename(path))[0] + '_add.mp4'))
def loaddata(video_index):
videoname = videonames[video_index]
img_index = random.randint(
int(N / 2) + 1, lengths[video_index] - int(N / 2) - 1)
input_img = np.zeros((opt.loadsize, opt.loadsize, 3 * N + 1),
dtype='uint8')
# this frame
this_mask = impro.imread(os.path.join(opt.dataset, videoname, 'mask',
'%05d' % (img_index) + '.png'),
'gray',
loadsize=opt.loadsize)
input_img[:, :, -1] = this_mask
#print(os.path.join(opt.dataset,videoname,'origin_image','%05d'%(img_index)+'.jpg'))
ground_true = impro.imread(os.path.join(opt.dataset, videoname,
'origin_image',
'%05d' % (img_index) + '.jpg'),
loadsize=opt.loadsize)
mosaic_size, mod, rect_rat, father = mosaic.get_random_parameter(
ground_true, this_mask)
# merge other frame
for i in range(0, N):
img = impro.imread(os.path.join(
opt.dataset, videoname, 'origin_image',
'%05d' % (img_index + i - int(N / 2)) + '.jpg'),
loadsize=opt.loadsize)
mask = impro.imread(os.path.join(
opt.dataset, videoname, 'mask',
'%05d' % (img_index + i - int(N / 2)) + '.png'),
'gray',
loadsize=opt.loadsize)
img_mosaic = mosaic.addmosaic_base(img,
mask,
mosaic_size,
model=mod,
rect_rat=rect_rat,
father=father)
input_img[:, :, i * 3:(i + 1) * 3] = img_mosaic
# to tensor
input_img, ground_true = data.random_transform_video(
input_img, ground_true, opt.finesize, N)
input_img = data.im2tensor(input_img,
bgr2rgb=False,
use_gpu=opt.use_gpu,
use_transform=False,
is0_1=False)
ground_true = data.im2tensor(ground_true,
bgr2rgb=False,
use_gpu=opt.use_gpu,
use_transform=False,
is0_1=False)
return input_img, ground_true
def cleanmosaic_video_fusion(opt,netG,netM):
path = opt.media_path
N = 25
INPUT_SIZE = 128
fps,imagepaths = video_init(opt,path)
positions = []
# get position
for i,imagepath in enumerate(imagepaths,1):
img_origin = impro.imread(os.path.join('./tmp/video2image',imagepath))
# x,y,size = runmodel.get_mosaic_position(img_origin,net_mosaic_pos,opt)[:3]
x,y,size,mask = runmodel.get_mosaic_position(img_origin,netM,opt)
cv2.imwrite(os.path.join('./tmp/mosaic_mask',imagepath), mask)
positions.append([x,y,size])
print('\r','Find mosaic location:'+str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='')
print('\nOptimize mosaic locations...')
positions =np.array(positions)
for i in range(3):positions[:,i] = filt.medfilt(positions[:,i],opt.medfilt_num)
# clean mosaic
for i,imagepath in enumerate(imagepaths,0):
x,y,size = positions[i][0],positions[i][1],positions[i][2]
img_origin = impro.imread(os.path.join('./tmp/video2image',imagepath))
mask = cv2.imread(os.path.join('./tmp/mosaic_mask',imagepath),0)
if size==0:
cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_origin)
else:
mosaic_input = np.zeros((INPUT_SIZE,INPUT_SIZE,3*N+1), dtype='uint8')
for j in range(0,N):
img = impro.imread(os.path.join('./tmp/video2image',imagepaths[np.clip(i+j-12,0,len(imagepaths)-1)]))
img = img[y-size:y+size,x-size:x+size]
img = impro.resize(img,INPUT_SIZE)
mosaic_input[:,:,j*3:(j+1)*3] = img
mask = impro.resize(mask,np.min(img_origin.shape[:2]))
mask = mask[y-size:y+size,x-size:x+size]
mask = impro.resize(mask, INPUT_SIZE)
mosaic_input[:,:,-1] = mask
mosaic_input = data.im2tensor(mosaic_input,bgr2rgb=False,use_gpu=opt.use_gpu,use_transform = False,is0_1 = False)
unmosaic_pred = netG(mosaic_input)
#unmosaic_pred = (unmosaic_pred.cpu().detach().numpy()*255)[0]
#img_fake = unmosaic_pred.transpose((1, 2, 0))
img_fake = data.tensor2im(unmosaic_pred,rgb2bgr = False ,is0_1 = False)
img_result = impro.replace_mosaic(img_origin,img_fake,x,y,size,opt.no_feather)
cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_result)
print('\r','Clean Mosaic:'+str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='')
print()
ffmpeg.image2video( fps,
'./tmp/replace_mosaic/output_%05d.'+opt.tempimage_type,
'./tmp/voice_tmp.mp3',
os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_clean.mp4'))
def loadimage(imagepaths,maskpaths,opt,test_flag = False):
batchsize = len(imagepaths)
images = np.zeros((batchsize,3,opt.finesize,opt.finesize), dtype=np.float32)
masks = np.zeros((batchsize,1,opt.finesize,opt.finesize), dtype=np.float32)
for i in range(len(imagepaths)):
img = impro.resize(impro.imread(imagepaths[i]),opt.loadsize)
mask = impro.resize(impro.imread(maskpaths[i],mod = 'gray'),opt.loadsize)
img,mask = data.random_transform_image(img, mask, opt.finesize, test_flag)
images[i] = (img.transpose((2, 0, 1))/255.0)
masks[i] = (mask.reshape(1,1,opt.finesize,opt.finesize)/255.0)
images = Totensor(images,opt.use_gpu)
masks = Totensor(masks,opt.use_gpu)
return images,masks
def get_mosaic_positions(opt, netM, imagepaths, savemask=True):
# get mosaic position
positions = []
t1 = time.time()
if not opt.no_preview:
cv2.namedWindow('mosaic mask', cv2.WINDOW_NORMAL)
print('Step:2/4 -- Find mosaic location')
for i, imagepath in enumerate(imagepaths, 1):
img_origin = impro.imread(
os.path.join(opt.temp_dir + '/video2image', imagepath))
x, y, size, mask = runmodel.get_mosaic_position(img_origin, netM, opt)
positions.append([x, y, size])
if savemask:
cv2.imwrite(os.path.join(opt.temp_dir + '/mosaic_mask', imagepath),
mask)
#preview result and print
if not opt.no_preview:
cv2.imshow('mosaic mask', mask)
cv2.waitKey(1) & 0xFF
t2 = time.time()
print('\r',
str(i) + '/' + str(len(imagepaths)),
util.get_bar(100 * i / len(imagepaths), num=35),
util.counttime(t1, t2, i, len(imagepaths)),
end='')
if not opt.no_preview:
cv2.destroyAllWindows()
print('\nOptimize mosaic locations...')
positions = np.array(positions)
for i in range(3):
positions[:, i] = filt.medfilt(positions[:, i], opt.medfilt_num)
return positions
def addmosaic_img(opt,netS):
path = opt.media_path
print('Add Mosaic:',path)
img = impro.imread(path)
mask = runmodel.get_ROI_position(img,netS,opt)[0]
img = mosaic.addmosaic(img,mask,opt)
impro.imwrite(os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_add.jpg'),img)
def cleanmosaic_video_byframe(opt, netG, netM):
path = opt.media_path
fps, imagepaths = video_init(opt, path)[:2]
positions = get_mosaic_positions(opt, netM, imagepaths, savemask=True)
# clean mosaic
for i, imagepath in enumerate(imagepaths, 0):
x, y, size = positions[i][0], positions[i][1], positions[i][2]
img_origin = impro.imread(os.path.join('./tmp/video2image', imagepath))
img_result = img_origin.copy()
if size != 0:
img_mosaic = img_origin[y - size:y + size, x - size:x + size]
if opt.traditional:
img_fake = runmodel.traditional_cleaner(img_mosaic, opt)
else:
img_fake = runmodel.run_pix2pix(img_mosaic, netG, opt)
mask = cv2.imread(os.path.join('./tmp/mosaic_mask', imagepath), 0)
img_result = impro.replace_mosaic(img_origin, img_fake, mask, x, y,
size, opt.no_feather)
cv2.imwrite(os.path.join('./tmp/replace_mosaic', imagepath),
img_result)
print('\r', 'Clean Mosaic:' + str(i + 1) + '/' + str(len(imagepaths)))
ffmpeg.image2video(
fps, './tmp/replace_mosaic/output_%05d.' + opt.tempimage_type,
'./tmp/voice_tmp.mp3',
os.path.join(
opt.result_dir,
os.path.splitext(os.path.basename(path))[0] + '_clean.mp4'))
def cleanmosaic_video_fusion(opt,netG,netM):
path = opt.media_path
N = 25
if 'HD' in os.path.basename(opt.model_path):
INPUT_SIZE = 256
else:
INPUT_SIZE = 128
fps,imagepaths,height,width = video_init(opt,path)
positions = get_mosaic_positions(opt,netM,imagepaths,savemask=True)
# clean mosaic
img_pool = np.zeros((height,width,3*N), dtype='uint8')
for i,imagepath in enumerate(imagepaths,0):
x,y,size = positions[i][0],positions[i][1],positions[i][2]
# image read stream
mask = cv2.imread(os.path.join('./tmp/mosaic_mask',imagepath),0)
if i==0 :
for j in range(0,N):
img_pool[:,:,j*3:(j+1)*3] = impro.imread(os.path.join('./tmp/video2image',imagepaths[np.clip(i+j-12,0,len(imagepaths)-1)]))
else:
img_pool[:,:,0:(N-1)*3] = img_pool[:,:,3:N*3]
img_pool[:,:,(N-1)*3:] = impro.imread(os.path.join('./tmp/video2image',imagepaths[np.clip(i+12,0,len(imagepaths)-1)]))
img_origin = img_pool[:,:,int((N-1)/2)*3:(int((N-1)/2)+1)*3]
if size==0: # can not find mosaic,
cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_origin)
else:
mosaic_input = np.zeros((INPUT_SIZE,INPUT_SIZE,3*N+1), dtype='uint8')
mosaic_input[:,:,0:N*3] = impro.resize(img_pool[y-size:y+size,x-size:x+size,:], INPUT_SIZE)
mask_input = impro.resize(mask,np.min(img_origin.shape[:2]))[y-size:y+size,x-size:x+size]
mosaic_input[:,:,-1] = impro.resize(mask_input, INPUT_SIZE)
mosaic_input = data.im2tensor(mosaic_input,bgr2rgb=False,use_gpu=opt.use_gpu,use_transform = False,is0_1 = False)
unmosaic_pred = netG(mosaic_input)
img_fake = data.tensor2im(unmosaic_pred,rgb2bgr = False ,is0_1 = False)
img_result = impro.replace_mosaic(img_origin,img_fake,mask,x,y,size,opt.no_feather)
cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_result)
print('\r','Clean Mosaic:'+str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='')
print()
ffmpeg.image2video( fps,
'./tmp/replace_mosaic/output_%05d.'+opt.tempimage_type,
'./tmp/voice_tmp.mp3',
os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_clean.mp4'))
def cleanmosaic_video_byframe(opt):
netG = loadmodel.pix2pix(opt)
net_mosaic_pos = loadmodel.unet_clean(opt)
path = opt.media_path
util.clean_tempfiles()
fps = ffmpeg.get_video_infos(path)[0]
ffmpeg.video2voice(path, './tmp/voice_tmp.mp3')
ffmpeg.video2image(path,
'./tmp/video2image/output_%05d.' + opt.tempimage_type)
positions = []
imagepaths = os.listdir('./tmp/video2image')
imagepaths.sort()
# get position
for imagepath in imagepaths:
img_origin = impro.imread(os.path.join('./tmp/video2image', imagepath))
x, y, size = runmodel.get_mosaic_position(img_origin, net_mosaic_pos,
opt)[:3]
positions.append([x, y, size])
print('Find mosaic location:', imagepath)
print('Optimize mosaic locations...')
positions = np.array(positions)
for i in range(3):
positions[:, i] = filt.medfilt(positions[:, i], opt.medfilt_num)
# clean mosaic
for i, imagepath in enumerate(imagepaths, 0):
x, y, size = positions[i][0], positions[i][1], positions[i][2]
img_origin = impro.imread(os.path.join('./tmp/video2image', imagepath))
img_result = img_origin.copy()
if size != 0:
img_mosaic = img_origin[y - size:y + size, x - size:x + size]
img_fake = runmodel.run_pix2pix(img_mosaic, netG, opt)
img_result = impro.replace_mosaic(img_origin, img_fake, x, y, size,
opt.no_feather)
cv2.imwrite(os.path.join('./tmp/replace_mosaic', imagepath),
img_result)
print('Clean Mosaic:', imagepath)
ffmpeg.image2video(
fps, './tmp/replace_mosaic/output_%05d.' + opt.tempimage_type,
'./tmp/voice_tmp.mp3',
os.path.join(
opt.result_dir,
os.path.splitext(os.path.basename(path))[0] + '_clean.mp4'))
def cleanmosaic_video_byframe(opt, netG, netM):
path = opt.media_path
fps, imagepaths = video_init(opt, path)[:2]
positions = get_mosaic_positions(opt, netM, imagepaths, savemask=True)
t1 = time.time()
if not opt.no_preview:
cv2.namedWindow('clean', cv2.WINDOW_NORMAL)
# clean mosaic
print('Step:3/4 -- Clean Mosaic:')
length = len(imagepaths)
for i, imagepath in enumerate(imagepaths, 0):
x, y, size = positions[i][0], positions[i][1], positions[i][2]
img_origin = impro.imread(
os.path.join(opt.temp_dir + '/video2image', imagepath))
img_result = img_origin.copy()
if size > 100:
try: #Avoid unknown errors
img_mosaic = img_origin[y - size:y + size, x - size:x + size]
if opt.traditional:
img_fake = runmodel.traditional_cleaner(img_mosaic, opt)
else:
img_fake = runmodel.run_pix2pix(img_mosaic, netG, opt)
mask = cv2.imread(
os.path.join(opt.temp_dir + '/mosaic_mask', imagepath), 0)
img_result = impro.replace_mosaic(img_origin, img_fake, mask,
x, y, size, opt.no_feather)
except Exception as e:
print('Warning:', e)
cv2.imwrite(os.path.join(opt.temp_dir + '/replace_mosaic', imagepath),
img_result)
os.remove(os.path.join(opt.temp_dir + '/video2image', imagepath))
#preview result and print
if not opt.no_preview:
cv2.imshow('clean', img_result)
cv2.waitKey(1) & 0xFF
t2 = time.time()
print('\r',
str(i + 1) + '/' + str(length),
util.get_bar(100 * i / length, num=35),
util.counttime(t1, t2, i + 1, len(imagepaths)),
end='')
print()
if not opt.no_preview:
cv2.destroyAllWindows()
print('Step:4/4 -- Convert images to video')
ffmpeg.image2video(
fps,
opt.temp_dir + '/replace_mosaic/output_%06d.' + opt.tempimage_type,
opt.temp_dir + '/voice_tmp.mp3',
os.path.join(
opt.result_dir,
os.path.splitext(os.path.basename(path))[0] + '_clean.mp4'))
def styletransfer_img(opt, netG):
print('Style Transfer_img:', opt.media_path)
img = impro.imread(opt.media_path)
img = runmodel.run_styletransfer(opt, netG, img)
suffix = os.path.basename(opt.model_path).replace('.pth', '').replace(
'style_', '')
impro.imwrite(
os.path.join(
opt.result_dir,
os.path.splitext(os.path.basename(opt.media_path))[0] + '_' +
suffix + '.jpg'), img)
def get_mosaic_positions(opt,netM,imagepaths,savemask=True):
# get mosaic position
positions = []
for i,imagepath in enumerate(imagepaths,1):
img_origin = impro.imread(os.path.join('./tmp/video2image',imagepath))
x,y,size,mask = runmodel.get_mosaic_position(img_origin,netM,opt)
if savemask:
cv2.imwrite(os.path.join('./tmp/mosaic_mask',imagepath), mask)
positions.append([x,y,size])
print('\r','Find mosaic location:'+str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='')
print('\nOptimize mosaic locations...')
positions =np.array(positions)
for i in range(3):positions[:,i] = filt.medfilt(positions[:,i],opt.medfilt_num)
return positions
def addmosaic_video(opt):
net = loadmodel.unet(opt)
path = opt.media_path
util.clean_tempfiles()
fps = ffmpeg.get_video_infos(path)[0]
ffmpeg.video2voice(path, './tmp/voice_tmp.mp3')
ffmpeg.video2image(path,
'./tmp/video2image/output_%05d.' + opt.tempimage_type)
imagepaths = os.listdir('./tmp/video2image')
imagepaths.sort()
# get position
positions = []
for imagepath in imagepaths:
print('Find ROI location:', imagepath)
img = impro.imread(os.path.join('./tmp/video2image', imagepath))
mask, x, y, area = runmodel.get_ROI_position(img, net, opt)
positions.append([x, y, area])
cv2.imwrite(os.path.join('./tmp/ROI_mask', imagepath), mask)
print('Optimize ROI locations...')
mask_index = filt.position_medfilt(np.array(positions), 7)
# add mosaic
print('Add mosaic to images...')
for i in range(len(imagepaths)):
mask = impro.imread(
os.path.join('./tmp/ROI_mask', imagepaths[mask_index[i]]))
img = impro.imread(os.path.join('./tmp/video2image', imagepaths[i]))
img = mosaic.addmosaic(img, mask, opt)
cv2.imwrite(os.path.join('./tmp/addmosaic_image', imagepaths[i]), img)
ffmpeg.image2video(
fps, './tmp/addmosaic_image/output_%05d.' + opt.tempimage_type,
'./tmp/voice_tmp.mp3',
os.path.join(opt.result_dir,
os.path.splitext(os.path.basename(path))[0] + '_add.mp4'))
def cleanmosaic_img(opt,netG,netM):
path = opt.media_path
print('Clean Mosaic:',path)
img_origin = impro.imread(path)
x,y,size,mask = runmodel.get_mosaic_position(img_origin,netM,opt)
#cv2.imwrite('./mask/'+os.path.basename(path), mask)
img_result = img_origin.copy()
if size != 0 :
img_mosaic = img_origin[y-size:y+size,x-size:x+size]
img_fake = runmodel.run_pix2pix(img_mosaic,netG,opt)
img_result = impro.replace_mosaic(img_origin,img_fake,x,y,size,opt.no_feather)
else:
print('Do not find mosaic')
impro.imwrite(os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_clean.jpg'),img_result)
def styletransfer_video(opt,netG):
path = opt.media_path
positions = []
fps,imagepaths = video_init(opt,path)
for i,imagepath in enumerate(imagepaths,1):
img = impro.imread(os.path.join('./tmp/video2image',imagepath))
img = runmodel.run_styletransfer(opt, netG, img)
cv2.imwrite(os.path.join('./tmp/style_transfer',imagepath),img)
print('\r','Transfer:'+str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='')
print()
suffix = os.path.basename(opt.model_path).replace('.pth','').replace('style_','')
ffmpeg.image2video( fps,
'./tmp/style_transfer/output_%05d.'+opt.tempimage_type,
'./tmp/voice_tmp.mp3',
os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_'+suffix+'.mp4'))
def styletransfer_video(opt, netG):
path = opt.media_path
positions = []
fps, imagepaths = video_init(opt, path)[:2]
print('Step:2/4 -- Transfer')
t1 = time.time()
if not opt.no_preview:
cv2.namedWindow('preview', cv2.WINDOW_NORMAL)
length = len(imagepaths)
for i, imagepath in enumerate(imagepaths, 1):
img = impro.imread(
os.path.join(opt.temp_dir + '/video2image', imagepath))
img = runmodel.run_styletransfer(opt, netG, img)
cv2.imwrite(os.path.join(opt.temp_dir + '/style_transfer', imagepath),
img)
os.remove(os.path.join(opt.temp_dir + '/video2image', imagepath))
#preview result and print
if not opt.no_preview:
cv2.imshow('preview', img)
cv2.waitKey(1) & 0xFF
t2 = time.time()
print('\r',
str(i) + '/' + str(length),
util.get_bar(100 * i / length, num=35),
util.counttime(t1, t2, i, len(imagepaths)),
end='')
print()
if not opt.no_preview:
cv2.destroyAllWindows()
suffix = os.path.basename(opt.model_path).replace('.pth', '').replace(
'style_', '')
print('Step:4/4 -- Convert images to video')
ffmpeg.image2video(
fps,
opt.temp_dir + '/style_transfer/output_%06d.' + opt.tempimage_type,
opt.temp_dir + '/voice_tmp.mp3',
os.path.join(
opt.result_dir,
os.path.splitext(os.path.basename(path))[0] + '_' + suffix +
'.mp4'))
timestamps = []
fps, endtime, height, width = ffmpeg.get_video_infos(videopath)
for cut_point in range(1, int(
(endtime - opt.time) / opt.interval)):
util.clean_tempfiles(opt)
ffmpeg.video2image(
videopath,
opt.temp_dir + '/video2image/%05d.' + opt.tempimage_type,
fps=1,
start_time=util.second2stamp(cut_point * opt.interval),
last_time=util.second2stamp(opt.time))
imagepaths = util.Traversal(opt.temp_dir + '/video2image')
imagepaths = sorted(imagepaths)
cnt = 0
for i in range(opt.time):
img = impro.imread(imagepaths[i])
mask = runmodel.get_ROI_position(img,
net,
opt,
keepsize=True)[0]
if not opt.all_mosaic_area:
mask = impro.find_mostlikely_ROI(mask)
x, y, size, area = impro.boundingSquare(mask, Ex_mul=1)
if area > opt.minmaskarea and size > opt.minsize and impro.Q_lapulase(
img) > opt.quality:
cnt += 1
if cnt == opt.time:
# print(second)
timestamps.append(
util.second2stamp(cut_point * opt.interval))
util.writelog(os.path.join(opt.savedir, 'opt.txt'),
print('network saved.')
#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)
def addmosaic_video(opt, netS):
path = opt.media_path
fps, imagepaths = video_init(opt, path)[:2]
length = len(imagepaths)
# get position
positions = []
t1 = time.time()
if not opt.no_preview:
cv2.namedWindow('preview', cv2.WINDOW_NORMAL)
print('Step:2/4 -- Find ROI location')
for i, imagepath in enumerate(imagepaths, 1):
img = impro.imread(
os.path.join(opt.temp_dir + '/video2image', imagepath))
mask, x, y, size, area = runmodel.get_ROI_position(img, netS, opt)
positions.append([x, y, area])
cv2.imwrite(os.path.join(opt.temp_dir + '/ROI_mask', imagepath), mask)
#preview result and print
if not opt.no_preview:
cv2.imshow('preview', mask)
cv2.waitKey(1) & 0xFF
t2 = time.time()
print('\r',
str(i) + '/' + str(length),
util.get_bar(100 * i / length, num=35),
util.counttime(t1, t2, i, length),
end='')
print('\nOptimize ROI locations...')
mask_index = filt.position_medfilt(np.array(positions), 7)
# add mosaic
print('Step:3/4 -- Add Mosaic:')
t1 = time.time()
for i, imagepath in enumerate(imagepaths, 1):
mask = impro.imread(
os.path.join(opt.temp_dir + '/ROI_mask',
imagepaths[mask_index[i - 1]]), 'gray')
img = impro.imread(
os.path.join(opt.temp_dir + '/video2image', imagepath))
if impro.mask_area(mask) > 100:
try: #Avoid unknown errors
img = mosaic.addmosaic(img, mask, opt)
except Exception as e:
print('Warning:', e)
cv2.imwrite(os.path.join(opt.temp_dir + '/addmosaic_image', imagepath),
img)
os.remove(os.path.join(opt.temp_dir + '/video2image', imagepath))
#preview result and print
if not opt.no_preview:
cv2.imshow('preview', img)
cv2.waitKey(1) & 0xFF
t2 = time.time()
print('\r',
str(i) + '/' + str(length),
util.get_bar(100 * i / length, num=35),
util.counttime(t1, t2, i, length),
end='')
print()
if not opt.no_preview:
cv2.destroyAllWindows()
print('Step:4/4 -- Convert images to video')
ffmpeg.image2video(
fps,
opt.temp_dir + '/addmosaic_image/output_%06d.' + opt.tempimage_type,
opt.temp_dir + '/voice_tmp.mp3',
os.path.join(opt.result_dir,
os.path.splitext(os.path.basename(path))[0] + '_add.mp4'))
def loader(imagepaths):
for imagepath in imagepaths:
img_origin = impro.imread(
os.path.join(opt.temp_dir + '/video2image', imagepath))
img_read_pool.put(img_origin)
irrpaths = util.Traversal(opt.irrholedir)
#def network
if 'network' in opt.mod:
net = loadmodel.bisenet(opt, 'roi')
print('Find images:', len(imgpaths))
starttime = datetime.datetime.now()
filecnt = 0
savecnt = opt.start
for fold in range(opt.fold):
for i in range(len(imgpaths)):
filecnt += 1
try:
# load image and get mask
img = impro.imread(imgpaths[i])
if 'drawn' in opt.mod:
mask_drawn = impro.imread(maskpaths[i], 'gray')
mask_drawn = impro.resize_like(mask_drawn, img)
mask = mask_drawn
if 'irregular' in opt.mod:
mask_irr = impro.imread(irrpaths[random.randint(0, 12000 - 1)],
'gray')
mask_irr = data.random_transform_single(
mask_irr, (img.shape[0], img.shape[1]))
mask = mask_irr
if 'network' in opt.mod:
mask_net = runmodel.get_ROI_position(img,
net,
opt,
keepsize=True)[0]
util.makedirs(train_path)
if MASK:
mask_path = os.path.join(output_dir, 'mask')
util.makedirs(mask_path)
mask_names = os.listdir(mask_dir)
img_names = os.listdir(img_dir)
mask_names.sort()
img_names.sort()
print('Find images:', len(img_names))
cnt = 0
for fold in range(FOLD_NUM):
for img_name, mask_name in zip(img_names, mask_names):
try:
img = impro.imread(os.path.join(img_dir, img_name))
mask = impro.imread(os.path.join(mask_dir, mask_name), 'gray')
mask = impro.resize_like(mask, img)
x, y, size, area = impro.boundingSquare(mask, 1.5)
if area > 100:
if Bounding:
img = impro.resize(
img[y - size:y + size, x - size:x + size], OUT_SIZE)
mask = impro.resize(
mask[y - size:y + size, x - size:x + size], OUT_SIZE)
img_mosaic = mosaic.addmosaic_random(img, mask)
if HD:
cv2.imwrite(
os.path.join(train_A_path, '%05d' % cnt + '.jpg'),
img_mosaic)
def cleanmosaic_video_fusion(opt):
net = loadmodel.pix2pix(opt)
net_mosaic_pos = loadmodel.unet_clean(opt)
path = opt.media_path
N = 25
util.clean_tempfiles()
fps = ffmpeg.get_video_infos(path)[0]
ffmpeg.video2voice(path, './tmp/voice_tmp.mp3')
ffmpeg.video2image(path,
'./tmp/video2image/output_%05d.' + opt.tempimage_type)
positions = []
imagepaths = os.listdir('./tmp/video2image')
imagepaths.sort()
# get position
for imagepath in imagepaths:
img_origin = impro.imread(os.path.join('./tmp/video2image', imagepath))
# x,y,size = runmodel.get_mosaic_position(img_origin,net_mosaic_pos,opt)[:3]
x, y, size, mask = runmodel.get_mosaic_position(
img_origin, net_mosaic_pos, opt)
cv2.imwrite(os.path.join('./tmp/mosaic_mask', imagepath), mask)
positions.append([x, y, size])
print('Find mosaic location:', imagepath)
print('Optimize mosaic locations...')
positions = np.array(positions)
for i in range(3):
positions[:, i] = filt.medfilt(positions[:, i], opt.medfilt_num)
# clean mosaic
print('Clean mosaic...')
for i, imagepath in enumerate(imagepaths, 0):
print('Clean mosaic:', imagepath)
x, y, size = positions[i][0], positions[i][1], positions[i][2]
img_origin = impro.imread(os.path.join('./tmp/video2image', imagepath))
mask = cv2.imread(os.path.join('./tmp/mosaic_mask', imagepath), 0)
if size == 0:
cv2.imwrite(os.path.join('./tmp/replace_mosaic', imagepath),
img_origin)
else:
mosaic_input = np.zeros((256, 256, 3 * N + 1), dtype='uint8')
for j in range(0, N):
img = impro.imread(
os.path.join(
'./tmp/video2image',
imagepaths[np.clip(i + j - 12, 0,
len(imagepaths) - 1)]))
img = img[y - size:y + size, x - size:x + size]
img = impro.resize(img, 256)
mosaic_input[:, :, j * 3:(j + 1) * 3] = img
mask = impro.resize(mask, np.min(img_origin.shape[:2]))
mask = mask[y - size:y + size, x - size:x + size]
mask = impro.resize(mask, 256)
mosaic_input[:, :, -1] = mask
mosaic_input = data.im2tensor(mosaic_input,
bgr2rgb=False,
use_gpu=opt.use_gpu,
use_transform=False)
unmosaic_pred = net(mosaic_input)
unmosaic_pred = (unmosaic_pred.cpu().detach().numpy() * 255)[0]
img_fake = unmosaic_pred.transpose((1, 2, 0))
img_result = impro.replace_mosaic(img_origin, img_fake, x, y, size,
opt.no_feather)
cv2.imwrite(os.path.join('./tmp/replace_mosaic', imagepath),
img_result)
ffmpeg.image2video(
fps, './tmp/replace_mosaic/output_%05d.' + opt.tempimage_type,
'./tmp/voice_tmp.mp3',
os.path.join(
opt.result_dir,
os.path.splitext(os.path.basename(path))[0] + '_clean.mp4'))
import numpy as np
import cv2
import os
import sys
sys.path.append("..")
from util import image_processing as impro
from util import util
img_dir = './datasets_img/pix2pix/edges2cat/images'
output_dir = './datasets_img/pix2pix/edges2cat/train'
util.makedirs(output_dir)
img_names = os.listdir(img_dir)
for i, img_name in enumerate(img_names, 2000):
try:
img = impro.imread(os.path.join(img_dir, img_name))
img = impro.resize(img, 286)
h, w = img.shape[:2]
edges = cv2.Canny(img, 150, 250)
edges = impro.ch_one2three(edges)
out_img = np.zeros((h, w * 2, 3), dtype=np.uint8)
out_img[:, 0:w] = edges
out_img[:, w:2 * w] = img
cv2.imwrite(os.path.join(output_dir, '%05d' % i + '.jpg'), out_img)
except Exception as e:
pass
mask = np.zeros(img_drawn.shape, np.uint8)
for row in range(img_drawn.shape[0]):
for col in range(img_drawn.shape[1]):
# if (img_drawn[row,col,:] == [0,255,0]).all(): #too slow
if img_drawn[row, col, 0] == 0:
if img_drawn[row, col, 1] == 255:
if img_drawn[row, col, 2] == 0:
mask[row, col, :] = [255, 255, 255]
return mask
cnt = 0
for file in filepaths:
try:
cnt += 1
img = impro.imread(file, loadsize=512)
img_drawn = img.copy()
cv2.namedWindow('image')
cv2.setMouseCallback('image', draw_circle) #MouseCallback
while (1):
cv2.imshow('image', img_drawn)
k = cv2.waitKey(1) & 0xFF
if k == ord('s'):
img_drawn = impro.resize(img_drawn, 256)
mask = makemask(img_drawn)
cv2.imwrite(
os.path.join(
mask_savedir,
os.path.splitext(os.path.basename(file))[0] + '.png'),
opt.use_gpu = True
net = loadmodel.unet(opt)
for path in videos:
path = os.path.join('./video', path)
util.clean_tempfiles()
ffmpeg.video2voice(path, './tmp/voice_tmp.mp3')
ffmpeg.video2image(path,
'./tmp/video2image/output_%05d.' + opt.tempimage_type)
imagepaths = os.listdir('./tmp/video2image')
imagepaths.sort()
# get position
positions = []
img_ori_example = impro.imread(
os.path.join('./tmp/video2image', imagepaths[0]))
mask_avg = np.zeros((impro.resize(img_ori_example, 128)).shape[:2])
for imagepath in imagepaths:
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)
def cleanmosaic_video_fusion(opt, netG, netM):
path = opt.media_path
N, T, S = 2, 5, 3
LEFT_FRAME = (N * S)
POOL_NUM = LEFT_FRAME * 2 + 1
INPUT_SIZE = 256
FRAME_POS = np.linspace(0, (T - 1) * S, T, dtype=np.int64)
img_pool = []
previous_frame = None
init_flag = True
fps, imagepaths, height, width = video_init(opt, path)
positions = get_mosaic_positions(opt, netM, imagepaths, savemask=True)
t1 = time.time()
if not opt.no_preview:
cv2.namedWindow('clean', cv2.WINDOW_NORMAL)
# clean mosaic
print('Step:3/4 -- Clean Mosaic:')
length = len(imagepaths)
for i, imagepath in enumerate(imagepaths, 0):
x, y, size = positions[i][0], positions[i][1], positions[i][2]
input_stream = []
# image read stream
if i == 0: # init
for j in range(POOL_NUM):
img_pool.append(
impro.imread(
os.path.join(
opt.temp_dir + '/video2image',
imagepaths[np.clip(i + j - LEFT_FRAME, 0,
len(imagepaths) - 1)])))
else: # load next frame
img_pool.pop(0)
img_pool.append(
impro.imread(
os.path.join(
opt.temp_dir + '/video2image',
imagepaths[np.clip(i + LEFT_FRAME, 0,
len(imagepaths) - 1)])))
img_origin = img_pool[LEFT_FRAME]
img_result = img_origin.copy()
if size > 50:
try: #Avoid unknown errors
for pos in FRAME_POS:
input_stream.append(
impro.resize(
img_pool[pos][y - size:y + size,
x - size:x + size],
INPUT_SIZE)[:, :, ::-1])
if init_flag:
init_flag = False
previous_frame = input_stream[N]
previous_frame = data.im2tensor(previous_frame,
bgr2rgb=True,
gpu_id=opt.gpu_id)
input_stream = np.array(input_stream).reshape(
1, T, INPUT_SIZE, INPUT_SIZE, 3).transpose((0, 4, 1, 2, 3))
input_stream = data.to_tensor(data.normalize(input_stream),
gpu_id=opt.gpu_id)
with torch.no_grad():
unmosaic_pred = netG(input_stream, previous_frame)
img_fake = data.tensor2im(unmosaic_pred, rgb2bgr=True)
previous_frame = unmosaic_pred
# previous_frame = data.tensor2im(unmosaic_pred,rgb2bgr = True)
mask = cv2.imread(
os.path.join(opt.temp_dir + '/mosaic_mask', imagepath), 0)
img_result = impro.replace_mosaic(img_origin, img_fake, mask,
x, y, size, opt.no_feather)
except Exception as e:
init_flag = True
print('Error:', e)
else:
init_flag = True
cv2.imwrite(os.path.join(opt.temp_dir + '/replace_mosaic', imagepath),
img_result)
os.remove(os.path.join(opt.temp_dir + '/video2image', imagepath))
#preview result and print
if not opt.no_preview:
cv2.imshow('clean', img_result)
cv2.waitKey(1) & 0xFF
t2 = time.time()
print('\r',
str(i + 1) + '/' + str(length),
util.get_bar(100 * i / length, num=35),
util.counttime(t1, t2, i + 1, len(imagepaths)),
end='')
print()
if not opt.no_preview:
cv2.destroyAllWindows()
print('Step:4/4 -- Convert images to video')
ffmpeg.image2video(
fps,
opt.temp_dir + '/replace_mosaic/output_%06d.' + opt.tempimage_type,
opt.temp_dir + '/voice_tmp.mp3',
os.path.join(
opt.result_dir,
os.path.splitext(os.path.basename(path))[0] + '_clean.mp4'))
