Skip to content

Commit

Permalink
Added an option to specify encoder codec in preprocess_video, reenact…
Browse files Browse the repository at this point in the history 
…, and swap scripts
  • Loading branch information
@YuvalNirkin
YuvalNirkin committed Apr 24, 2020
1 parent 05c9ec7 commit 5e61edb
Show file tree
Hide file tree
Showing 4 changed files with 20 additions and 14 deletions.
11 changes: 6 additions & 5 deletions inference/reenact.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@

class FaceReenactment(VideoProcessBase):
def __init__(self, resolution=d('resolution'), crop_scale=d('crop_scale'), gpus=d('gpus'),
cpu_only=d('cpu_only'), display=d('display'), verbose=d('verbose'),
cpu_only=d('cpu_only'), display=d('display'), verbose=d('verbose'), encoder_codec=d('encoder_codec'),
# Detection arguments:
detection_model=d('detection_model'), det_batch_size=d('det_batch_size'), det_postfix=d('det_postfix'),
# Sequence arguments:
Expand All @@ -97,7 +97,7 @@ def __init__(self, resolution=d('resolution'), crop_scale=d('crop_scale'), gpus=
batch_size=d('batch_size'), reenactment_model=d('reenactment_model'), criterion_id=d('criterion_id'),
min_radius=d('min_radius')):
super(FaceReenactment, self).__init__(
resolution, crop_scale, gpus, cpu_only, display, verbose,
resolution, crop_scale, gpus, cpu_only, display, verbose, encoder_codec,
detection_model=detection_model, det_batch_size=det_batch_size, det_postfix=det_postfix,
iou_thresh=iou_thresh, min_length=min_length, min_size=min_size, center_kernel=center_kernel,
size_kernel=size_kernel, smooth_det=smooth_det, seq_postfix=seq_postfix, write_empty=write_empty,
Expand Down Expand Up @@ -138,7 +138,7 @@ def __init__(self, resolution=d('resolution'), crop_scale=d('crop_scale'), gpus=

# Initialize video renderer
self.video_renderer = FaceReenactmentRenderer(self.display, self.verbose, True, self.resolution,
self.crop_scale)
self.crop_scale, encoder_codec)
self.video_renderer.start()

def __del__(self):
Expand Down Expand Up @@ -308,7 +308,8 @@ def __call__(self, source_path, target_path, output_path=None, select_source='lo


class FaceReenactmentRenderer(VideoRenderer):
def __init__(self, display=False, verbose=0, output_crop=False, resolution=256, crop_scale=1.2):
def __init__(self, display=False, verbose=0, output_crop=False, resolution=256, crop_scale=1.2,
encoder_codec='avc1'):
self._appearance_map = None
self._fig = None
self._figsize = (24, 16)
Expand All @@ -324,7 +325,7 @@ def __init__(self, display=False, verbose=0, output_crop=False, resolution=256,
verbose_size = (self._appearance_map_size[0] + resolution, self._appearance_map_size[1])

super(FaceReenactmentRenderer, self).__init__(display, verbose, verbose_size, output_crop, resolution,
crop_scale)
crop_scale, encoder_codec)

def on_render(self, *args):
if self._verbose <= 0:
Expand Down
10 changes: 6 additions & 4 deletions inference/swap.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@

class FaceSwapping(VideoProcessBase):
def __init__(self, resolution=d('resolution'), crop_scale=d('crop_scale'), gpus=d('gpus'),
cpu_only=d('cpu_only'), display=d('display'), verbose=d('verbose'),
cpu_only=d('cpu_only'), display=d('display'), verbose=d('verbose'), encoder_codec=d('encoder_codec'),
# Detection arguments:
detection_model=d('detection_model'), det_batch_size=d('det_batch_size'), det_postfix=d('det_postfix'),
# Sequence arguments:
Expand Down Expand Up @@ -157,7 +157,7 @@ def __init__(self, resolution=d('resolution'), crop_scale=d('crop_scale'), gpus=

# Initialize video writer
self.video_renderer = FaceSwappingRenderer(self.display, self.verbose, self.output_crop, self.resolution,
self.crop_scale)
self.crop_scale, encoder_codec)
self.video_renderer.start()

def __del__(self):
Expand Down Expand Up @@ -358,7 +358,8 @@ def __call__(self, source_path, target_path, output_path=None, select_source='lo


class FaceSwappingRenderer(VideoRenderer):
def __init__(self, display=False, verbose=0, output_crop=False, resolution=256, crop_scale=1.2):
def __init__(self, display=False, verbose=0, output_crop=False, resolution=256, crop_scale=1.2,
encoder_codec='avc1'):
self._appearance_map = None
self._fig = None
self._figsize = (24, 16)
Expand All @@ -373,7 +374,8 @@ def __init__(self, display=False, verbose=0, output_crop=False, resolution=256,
self._appearance_map_size = (int(np.round(height * fig_ratio)), height)
verbose_size = (self._appearance_map_size[0] + resolution * 2, self._appearance_map_size[1])

super(FaceSwappingRenderer, self).__init__(display, verbose, verbose_size, output_crop, resolution, crop_scale)
super(FaceSwappingRenderer, self).__init__(display, verbose, verbose_size, output_crop, resolution, crop_scale,
encoder_codec)

def on_render(self, *args):
if self._verbose <= 0:
Expand Down
7 changes: 4 additions & 3 deletions preprocess/preprocess_video.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,8 @@
help='display the rendering')
general.add_argument('-v', '--verbose', default=0, type=int, metavar='N',
help='verbose level')
general.add_argument('-ec', '--encoder_codec', default='avc1', metavar='STR',
help='encoder codec code')

detection = base_parser.add_argument_group('detection')
detection.add_argument('-dm', '--detection_model', metavar='PATH', default='../weights/WIDERFace_DSFD_RES152.pth',
Expand Down Expand Up @@ -124,7 +126,7 @@

class VideoProcessBase(object):
def __init__(self, resolution=d('resolution'), crop_scale=d('crop_scale'), gpus=d('gpus'),
cpu_only=d('cpu_only'), display=d('display'), verbose=d('verbose'),
cpu_only=d('cpu_only'), display=d('display'), verbose=d('verbose'), encoder_codec=d('encoder_codec'),
# Detection arguments:
detection_model=d('detection_model'), det_batch_size=d('det_batch_size'), det_postfix=d('det_postfix'),
# Sequence arguments:
Expand Down Expand Up @@ -213,7 +215,7 @@ def __init__(self, resolution=d('resolution'), crop_scale=d('crop_scale'), gpus=
self.smooth_seg = None

# Initialize output videos format
self.fourcc = cv2.VideoWriter_fourcc(*'avc1')
self.fourcc = cv2.VideoWriter_fourcc(*encoder_codec)

def process_pose(self, input_path, output_dir, seq_file_path):
if not self.cache_pose:
Expand Down Expand Up @@ -538,4 +540,3 @@ def main(input, output=d('output'), resolution=d('resolution'), crop_scale=d('cr

if __name__ == "__main__":
main(**vars(parser.parse_args()))

6 changes: 4 additions & 2 deletions utils/video_renderer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,10 @@ class VideoRenderer(mp.Process):
verbose level zero
resolution (int): Determines the size of cropped frames to be (resolution, resolution)
crop_scale (float): Multiplier factor to scale tight bounding boxes
encoder_codec (str) Encoder codec code
"""
def __init__(self, display=False, verbose=0, verbose_size=None, output_crop=False, resolution=256, crop_scale=1.2):
def __init__(self, display=False, verbose=0, verbose_size=None, output_crop=False, resolution=256, crop_scale=1.2,
encoder_codec='avc1'):
super(VideoRenderer, self).__init__()
self._display = display
self._verbose = verbose
Expand All @@ -32,7 +34,7 @@ def __init__(self, display=False, verbose=0, verbose_size=None, output_crop=Fals
self._running = True
self._input_queue = mp.Queue()
self._reply_queue = mp.Queue()
self._fourcc = cv2.VideoWriter_fourcc(*'avc1')
self._fourcc = cv2.VideoWriter_fourcc(*'encoder_codec')
self._in_vid = None
self._out_vid = None
self._seq = None
Expand Down

0 comments on commit 5e61edb

Please sign in to comment.