Merge pull request #23 from LielinJiang/sr

Move some model to ppgan, Add sr model
AnnaTrainingG · Sep 18, 2020 · 8a4848d · 8a4848d
2 parents 3bc13ff + d5b9c2e
commit 8a4848d
Show file tree

Hide file tree

Showing 33 changed files with 891 additions and 326 deletions.
diff --git a/applications/DAIN/predict.py b/applications/DAIN/predict.py
@@ -13,8 +13,8 @@
 
 import paddle.fluid as fluid
 from paddle.utils.download import get_path_from_url
+from ppgan.utils.video import video2frames, frames2video
 
-import networks
 from util import *
 from my_args import parser
 
@@ -129,7 +129,7 @@ def run(self):
             r2 = str(int(fps) * times_interp)
             print("New fps (frame rate): ", r2)
 
-            out_path = dump_frames_ffmpeg(vid, frame_path_input)
+            out_path = video2frames(vid, frame_path_input)
 
             vidname = vid.split('/')[-1].split('.')[0]
 
@@ -266,7 +266,7 @@ def run(self):
                                                 vidname + '.mp4')
             if os.path.exists(video_pattern_output):
                 os.remove(video_pattern_output)
-            frames_to_video_ffmpeg(frame_pattern_combined, video_pattern_output,
+            frames2video(frame_pattern_combined, video_pattern_output,
                                    r2)
 
         return frame_pattern_combined, video_pattern_output

diff --git a/applications/DAIN/util.py b/applications/DAIN/util.py
@@ -21,66 +21,6 @@ def update(self, val, n=1):
         self.avg = self.sum / self.count
 
 
-def dump_frames_ffmpeg(vid_path, outpath, r=None, ss=None, t=None):
-    ffmpeg = ['ffmpeg ', ' -y -loglevel ', ' error ']
-    vid_name = vid_path.split('/')[-1].split('.')[0]
-    out_full_path = os.path.join(outpath, vid_name)
-
-    if not os.path.exists(out_full_path):
-        os.makedirs(out_full_path)
-
-    # video file name
-    outformat = out_full_path + '/%08d.png'
-
-    if ss is not None and t is not None and r is not None:
-        cmd = ffmpeg + [
-            ' -ss ',
-            ss,
-            ' -t ',
-            t,
-            ' -i ',
-            vid_path,
-            ' -r ',
-            r,
-            # ' -f ', ' image2 ',
-            #                        ' -s ', ' 960*540 ',
-            ' -qscale:v ',
-            ' 0.1 ',
-            ' -start_number ',
-            ' 0 ',
-            # ' -qmax ', ' 1 ',
-            outformat
-        ]
-    else:
-        cmd = ffmpeg + [' -i ', vid_path, ' -start_number ', ' 0 ', outformat]
-
-    cmd = ''.join(cmd)
-
-    if os.system(cmd) == 0:
-        pass
-    else:
-        print('ffmpeg process video: {} error'.format(vid_name))
-
-    sys.stdout.flush()
-    return out_full_path
-
-
-def frames_to_video_ffmpeg(framepath, videopath, r):
-    ffmpeg = ['ffmpeg ', ' -y -loglevel ', ' error ']
-    cmd = ffmpeg + [
-        ' -r ', r, ' -f ', ' image2 ', ' -i ', framepath, ' -vcodec ',
-        ' libx264 ', ' -pix_fmt ', ' yuv420p ', ' -crf ', ' 16 ', videopath
-    ]
-    cmd = ''.join(cmd)
-
-    if os.system(cmd) == 0:
-        pass
-    else:
-        print('ffmpeg process video: {} error'.format(videopath))
-
-    sys.stdout.flush()
-
-
 def combine_frames(input, interpolated, combined, num_frames):
     frames1 = sorted(glob.glob(os.path.join(input, '*.png')))
     frames2 = sorted(glob.glob(os.path.join(interpolated, '*.png')))

diff --git a/applications/DeOldify/predict.py b/applications/DeOldify/predict.py
@@ -14,8 +14,9 @@
 from PIL import Image
 from tqdm import tqdm
 from paddle import fluid
-from model import build_model
 from paddle.utils.download import get_path_from_url
+from ppgan.utils.video import frames2video, video2frames
+from ppgan.models.generators.deoldify import build_model
 
 parser = argparse.ArgumentParser(description='DeOldify')
 parser.add_argument('--input', type=str, default='none', help='Input video')
@@ -29,23 +30,7 @@
                     default=None,
                     help='Path to the reference image directory')
 
-DeOldify_weight_url = 'https://paddlegan.bj.bcebos.com/applications/DeOldify_stable.pdparams'
-
-
-def frames_to_video_ffmpeg(framepath, videopath, r):
-    ffmpeg = ['ffmpeg ', ' -y -loglevel ', ' error ']
-    cmd = ffmpeg + [
-        ' -r ', r, ' -f ', ' image2 ', ' -i ', framepath, ' -vcodec ',
-        ' libx264 ', ' -pix_fmt ', ' yuv420p ', ' -crf ', ' 16 ', videopath
-    ]
-    cmd = ''.join(cmd)
-
-    if os.system(cmd) == 0:
-        pass
-    else:
-        print('ffmpeg process video: {} error'.format(videopath))
-
-    sys.stdout.flush()
+DEOLDIFY_WEIGHT_URL = 'https://paddlegan.bj.bcebos.com/applications/DeOldify_stable.pdparams'
 
 
 class DeOldifyPredictor():
@@ -60,7 +45,7 @@ def __init__(self,
         self.render_factor = render_factor
         self.model = build_model()
         if weight_path is None:
-            weight_path = get_path_from_url(DeOldify_weight_url, cur_path)
+            weight_path = get_path_from_url(DEOLDIFY_WEIGHT_URL, cur_path)
 
         state_dict, _ = paddle.load(weight_path)
         self.model.load_dict(state_dict)
@@ -127,7 +112,7 @@ def run(self):
         cap = cv2.VideoCapture(vid)
         fps = cap.get(cv2.CAP_PROP_FPS)
 
-        out_path = dump_frames_ffmpeg(vid, output_path)
+        out_path = video2frames(vid, output_path)
 
         frames = sorted(glob.glob(os.path.join(out_path, '*.png')))
 
@@ -141,42 +126,11 @@ def run(self):
 
         vid_out_path = os.path.join(output_path,
                                     '{}_deoldify_out.mp4'.format(base_name))
-        frames_to_video_ffmpeg(frame_pattern_combined, vid_out_path,
-                               str(int(fps)))
+        frames2video(frame_pattern_combined, vid_out_path, str(int(fps)))
 
         return frame_pattern_combined, vid_out_path
 
 
-def dump_frames_ffmpeg(vid_path, outpath, r=None, ss=None, t=None):
-    ffmpeg = ['ffmpeg ', ' -y -loglevel ', ' error ']
-    vid_name = vid_path.split('/')[-1].split('.')[0]
-    out_full_path = os.path.join(outpath, 'frames_input')
-
-    if not os.path.exists(out_full_path):
-        os.makedirs(out_full_path)
-
-    # video file name
-    outformat = out_full_path + '/%08d.png'
-
-    if ss is not None and t is not None and r is not None:
-        cmd = ffmpeg + [
-            ' -ss ', ss, ' -t ', t, ' -i ', vid_path, ' -r ', r, ' -qscale:v ',
-            ' 0.1 ', ' -start_number ', ' 0 ', outformat
-        ]
-    else:
-        cmd = ffmpeg + [' -i ', vid_path, ' -start_number ', ' 0 ', outformat]
-
-    cmd = ''.join(cmd)
-
-    if os.system(cmd) == 0:
-        pass
-    else:
-        print('ffmpeg process video: {} error'.format(vid_name))
-
-    sys.stdout.flush()
-    return out_full_path
-
-
 if __name__ == '__main__':
     paddle.disable_static()
     args = parser.parse_args()

diff --git a/applications/DeOldify/spectral_norm.py b/applications/DeOldify/spectral_norm.py
diff --git a/applications/DeepRemaster/predict.py b/applications/DeepRemaster/predict.py
@@ -14,10 +14,10 @@
 import argparse
 import subprocess
 import utils
-from remasternet import NetworkR, NetworkC
+from ppgan.models.generators.remaster import NetworkR, NetworkC
 from paddle.utils.download import get_path_from_url
 
-DeepRemaster_weight_url = 'https://paddlegan.bj.bcebos.com/applications/deep_remaster.pdparams'
+DEEPREMASTER_WEIGHT_URL = 'https://paddlegan.bj.bcebos.com/applications/deep_remaster.pdparams'
 
 parser = argparse.ArgumentParser(description='Remastering')
 parser.add_argument('--input', type=str, default=None, help='Input video')
@@ -51,7 +51,7 @@ def __init__(self,
         self.mindim = mindim
 
         if weight_path is None:
-            weight_path = get_path_from_url(DeepRemaster_weight_url, cur_path)
+            weight_path = get_path_from_url(DEEPREMASTER_WEIGHT_URL, cur_path)
 
         state_dict, _ = paddle.load(weight_path)
 

diff --git a/applications/EDVR/predict.py b/applications/EDVR/predict.py
@@ -30,8 +30,9 @@
 from tqdm import tqdm
 from data import EDVRDataset
 from paddle.utils.download import get_path_from_url
+from ppgan.utils.video import frames2video, video2frames
 
-EDVR_weight_url = 'https://paddlegan.bj.bcebos.com/applications/edvr_infer_model.tar'
+EDVR_WEIGHT_URL = 'https://paddlegan.bj.bcebos.com/applications/edvr_infer_model.tar'
 
 
 def parse_args():
@@ -71,52 +72,6 @@ def save_img(img, framename):
     cv2.imwrite(framename, img)
 
 
-def dump_frames_ffmpeg(vid_path, outpath, r=None, ss=None, t=None):
-    ffmpeg = ['ffmpeg ', ' -y -loglevel ', ' error ']
-    vid_name = vid_path.split('/')[-1].split('.')[0]
-    out_full_path = os.path.join(outpath, 'frames_input')
-
-    if not os.path.exists(out_full_path):
-        os.makedirs(out_full_path)
-
-    # video file name
-    outformat = out_full_path + '/%08d.png'
-
-    if ss is not None and t is not None and r is not None:
-        cmd = ffmpeg + [
-            ' -ss ', ss, ' -t ', t, ' -i ', vid_path, ' -r ', r, ' -qscale:v ',
-            ' 0.1 ', ' -start_number ', ' 0 ', outformat
-        ]
-    else:
-        cmd = ffmpeg + [' -i ', vid_path, ' -start_number ', ' 0 ', outformat]
-
-    cmd = ''.join(cmd)
-
-    if os.system(cmd) == 0:
-        pass
-    else:
-        print('ffmpeg process video: {} error'.format(vid_name))
-
-    sys.stdout.flush()
-    return out_full_path
-
-
-def frames_to_video_ffmpeg(framepath, videopath, r):
-    ffmpeg = ['ffmpeg ', ' -y -loglevel ', ' error ']
-    cmd = ffmpeg + [
-        ' -r ', r, ' -f ', ' image2 ', ' -i ', framepath, ' -vcodec ',
-        ' libx264 ', ' -pix_fmt ', ' yuv420p ', ' -crf ', ' 16 ', videopath
-    ]
-    cmd = ''.join(cmd)
-
-    if os.system(cmd) == 0:
-        pass
-    else:
-        print('ffmpeg process video: {} error'.format(videopath))
-
-    sys.stdout.flush()
-
-
 class EDVRPredictor:
     def __init__(self, input, output, weight_path=None):
         self.input = input
@@ -127,9 +82,7 @@ def __init__(self, input, output, weight_path=None):
         self.exe = fluid.Executor(place)
 
         if weight_path is None:
-            weight_path = get_path_from_url(EDVR_weight_url, cur_path)
-
-        print(weight_path)
+            weight_path = get_path_from_url(EDVR_WEIGHT_URL, cur_path)
 
         model_filename = 'EDVR_model.pdmodel'
         params_filename = 'EDVR_params.pdparams'
@@ -155,7 +108,7 @@ def run(self):
         cap = cv2.VideoCapture(vid)
         fps = cap.get(cv2.CAP_PROP_FPS)
 
-        out_path = dump_frames_ffmpeg(vid, output_path)
+        out_path = video2frames(vid, output_path)
 
         frames = sorted(glob.glob(os.path.join(out_path, '*.png')))
 
@@ -188,8 +141,7 @@ def run(self):
         frame_pattern_combined = os.path.join(pred_frame_path, '%08d.png')
         vid_out_path = os.path.join(self.output,
                                     '{}_edvr_out.mp4'.format(base_name))
-        frames_to_video_ffmpeg(frame_pattern_combined, vid_out_path,
-                               str(int(fps)))
+        frames2video(frame_pattern_combined, vid_out_path, str(int(fps)))
 
         return frame_pattern_combined, vid_out_path