def addmosaic_video(opt, netS):
path = opt.media_path
fps, imagepaths = video_init(opt, path)[:2]
# get position
positions = []
for i, imagepath in enumerate(imagepaths, 1):
img = impro.imread(os.path.join('./tmp/video2image', imagepath))
mask, x, y, size, 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=35),
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=35),
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 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)),
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.results_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 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 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 __mosaic(opt, A_path, B_path):
ground_true = impro.imread(B_path)
mask = impro.imread(A_path, 'gray')
mosaic_size, mod, rect_rat, feather = mosaic.get_random_parameter(
ground_true, mask)
img_mosaic = mosaic.addmosaic_base(ground_true,
mask,
mosaic_size,
model=mod,
rect_rat=rect_rat,
feather=-1)
image_mosaic = Image.fromarray(cv2.cvtColor(img_mosaic, cv2.COLOR_BGR2RGB))
dir_A_sample = os.path.join(opt.dataroot, opt.phase + "_A_sample")
if os.path.isdir(dir_A_sample):
r = random.randint(1, 1000)
if r == 1:
image_mosaic.save(
os.path.join(
dir_A_sample,
str(time.time()) + "_" + str(mosaic_size) + ".jpg"))
return image_mosaic
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 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]
if opt.traditional:
img_fake = runmodel.traditional_cleaner(img_mosaic, opt)
else:
img_fake = runmodel.run_pix2pix(img_mosaic, netG, opt)
img_result = impro.replace_mosaic(img_origin, img_fake, mask, 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)[:2]
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=35),
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 cleanmosaic_video_fusion(opt, netG, netM):
path = opt.media_path
N = opt.N
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'))
# generate datasets
print('Generate datasets...')
test_B_path = os.path.join(opt.savedir, 'test_B')
test_A_path = os.path.join(opt.savedir, 'test_A')
util.makedirs(test_B_path)
util.makedirs(test_A_path)
imagepaths = util.Traversal(opt.datadir)
imagepaths = sorted(imagepaths)
imgs = []
masks = []
mask_flag = False
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))