Skip to content

Commit

Permalink
1. Allow unfinished tasks to be restored (#38 #33)
Browse files Browse the repository at this point in the history 
2. Optimize image_pool when processing video
  • Loading branch information
@HypoX64
HypoX64 committed Dec 17, 2020
1 parent 6ca8404 commit cab98f7
Show file tree
Hide file tree
Showing 4 changed files with 54 additions and 27 deletions.
51 changes: 35 additions & 16 deletions cores/core.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,14 +11,32 @@
---------------------Video Init---------------------
'''
def video_init(opt,path):
util.clean_tempfiles(opt)
fps,endtime,height,width = ffmpeg.get_video_infos(path)
if opt.fps !=0:
fps = opt.fps
ffmpeg.video2voice(path,opt.temp_dir+'/voice_tmp.mp3',opt.start_time,opt.last_time)
ffmpeg.video2image(path,opt.temp_dir+'/video2image/output_%06d.'+opt.tempimage_type,fps,opt.start_time,opt.last_time)
imagepaths=os.listdir(opt.temp_dir+'/video2image')
imagepaths.sort()

continue_flag = False
if os.path.isdir(opt.temp_dir):
if (opt.last_time != '00:00:00' and len(os.listdir(os.path.join(opt.temp_dir,'video2image'))) > fps*(util.stamp2second(opt.last_time)-1)) \
or (opt.last_time == '00:00:00' and len(os.listdir(os.path.join(opt.temp_dir,'video2image'))) > fps*(endtime-1)):
choose = input('There is an unfinished video. Continue it? [y/n] ')
if choose.lower() =='yes' or choose.lower() == 'y':
continue_flag = True

if continue_flag:
processed_num = max(len(os.listdir(os.path.join(opt.temp_dir,'addmosaic_image'))),
len(os.listdir(os.path.join(opt.temp_dir,'replace_mosaic'))),
len(os.listdir(os.path.join(opt.temp_dir,'style_transfer'))))
imagepaths = os.listdir(opt.temp_dir+'/video2image')
imagepaths.sort()
imagepaths = imagepaths[processed_num:]
else:
util.file_init(opt)
ffmpeg.video2voice(path,opt.temp_dir+'/voice_tmp.mp3',opt.start_time,opt.last_time)
ffmpeg.video2image(path,opt.temp_dir+'/video2image/output_%06d.'+opt.tempimage_type,fps,opt.start_time,opt.last_time)
imagepaths = os.listdir(opt.temp_dir+'/video2image')
imagepaths.sort()

return fps,imagepaths,height,width

'''
Expand Down Expand Up @@ -238,34 +256,35 @@ def cleanmosaic_video_fusion(opt,netG,netM):
# clean mosaic
print('Clean Mosaic:')
length = len(imagepaths)
img_pool = np.zeros((height,width,3*N), dtype='uint8')

img_pool = []
mosaic_input = np.zeros((INPUT_SIZE,INPUT_SIZE,3*N+1), 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(opt.temp_dir+'/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(opt.temp_dir+'/video2image',imagepaths[np.clip(i+j-12,0,len(imagepaths)-1)]))
img_pool.append(impro.imread(os.path.join(opt.temp_dir+'/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(opt.temp_dir+'/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]
img_pool.pop(0)
img_pool.append(impro.imread(os.path.join(opt.temp_dir+'/video2image',imagepaths[np.clip(i+12,0,len(imagepaths)-1)])))
img_origin = img_pool[12]
img_result = img_origin.copy()

if size>100:
try:#Avoid unknown errors
#reshape to network input shape

mosaic_input[:,:,0:N*3] = impro.resize(img_pool[y-size:y+size,x-size:x+size,:], INPUT_SIZE)
for k in range(N):
mosaic_input[:,:,k*3:(k+1)*3] = impro.resize(img_pool[k][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_tensor = data.im2tensor(mosaic_input,bgr2rgb=False,use_gpu=opt.use_gpu,use_transform = False,is0_1 = False)
unmosaic_pred = netG(mosaic_input_tensor)
unmosaic_pred = netG(mosaic_input_tensor)
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)
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)
Expand Down
7 changes: 4 additions & 3 deletions cores/options.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import argparse
import os
import sys


class Options():
Expand Down Expand Up @@ -69,7 +70,7 @@ def getparse(self, test_flag = False):
if not os.path.exists(self.opt.media_path):
print('Error: Bad media path!')
input('Please press any key to exit.\n')
exit(0)
sys.exit(0)

if self.opt.mode == 'auto':
if 'clean' in model_name or self.opt.traditional:
Expand All @@ -81,7 +82,7 @@ def getparse(self, test_flag = False):
else:
print('Please input running model!')
input('Please press any key to exit.\n')
exit(0)
sys.exit(0)

if self.opt.output_size == 0 and self.opt.mode == 'style':
self.opt.output_size = 512
Expand All @@ -101,7 +102,7 @@ def getparse(self, test_flag = False):
else:
print('Type of Generator error!')
input('Please press any key to exit.\n')
exit(0)
sys.exit(0)

if self.opt.ex_mult == 'auto':
if 'face' in model_name:
Expand Down
10 changes: 7 additions & 3 deletions deepmosaic.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,8 @@
from models import loadmodel

opt = Options().getparse(test_flag = True)
util.file_init(opt)
if not os.path.isdir(opt.temp_dir):
util.file_init(opt)

def main():

Expand All @@ -22,8 +23,10 @@ def main():
core.addmosaic_img(opt,netS)
elif util.is_video(file):
core.addmosaic_video(opt,netS)
util.clean_tempfiles(opt, tmp_init = False)
else:
print('This type of file is not supported')
util.clean_tempfiles(opt, tmp_init = False)

elif opt.mode == 'clean':
netM = loadmodel.bisenet(opt,'mosaic')
Expand All @@ -43,6 +46,7 @@ def main():
core.cleanmosaic_video_fusion(opt,netG,netM)
else:
core.cleanmosaic_video_byframe(opt,netG,netM)
util.clean_tempfiles(opt, tmp_init = False)
else:
print('This type of file is not supported')

Expand All @@ -54,6 +58,7 @@ def main():
core.styletransfer_img(opt,netG)
elif util.is_video(file):
core.styletransfer_video(opt,netG)
util.clean_tempfiles(opt, tmp_init = False)
else:
print('This type of file is not supported')

Expand Down Expand Up @@ -83,5 +88,4 @@ def main():
print(stack)
input('Please press any key to exit.\n')
#util.clean_tempfiles(tmp_init = False)
exit(0)

sys.exit(0)
13 changes: 8 additions & 5 deletions util/util.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,11 +70,9 @@ def clean_tempfiles(opt,tmp_init=True):
os.makedirs(tmpdir)
os.makedirs(os.path.join(tmpdir, 'video2image'))
os.makedirs(os.path.join(tmpdir, 'addmosaic_image'))
os.makedirs(os.path.join(tmpdir, 'mosaic_crop'))
os.makedirs(os.path.join(tmpdir, 'replace_mosaic'))
os.makedirs(os.path.join(tmpdir, 'mosaic_mask'))
os.makedirs(os.path.join(tmpdir, 'ROI_mask'))
os.makedirs(os.path.join(tmpdir, 'ROI_mask_check'))
os.makedirs(os.path.join(tmpdir, 'style_transfer'))

def file_init(opt):
Expand All @@ -90,11 +88,16 @@ def second2stamp(s):
s = int(s%60)
return "%02d:%02d:%02d" % (h, m, s)

def counttime(t1,t2,now_num,all_num):
def stamp2second(stamp):
substamps = stamp.split(':')
return int(substamps[0])*3600 + int(substamps[1])*60 + int(substamps[2])


def counttime(start_time,current_time,now_num,all_num):
'''
t1,t2: time.time()
start_time,current_time: time.time()
'''
used_time = int(t2-t1)
used_time = int(current_time-start_time)
all_time = int(used_time/now_num*all_num)
return second2stamp(used_time)+'/'+second2stamp(all_time)

Expand Down

2 comments on commit cab98f7

@deepmos01
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice man. Should I close my issue? This commit solves it?

@HypoX64
Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I only updated the source code. The exe version does not currently have this feature, and I will update it later.

Please sign in to comment.