[Feature] Support different filename templates in GenerateSegmentIndices

demo/restoration_video_demo.py

@@ -13,6 +13,10 @@ def parse_args():
     parser.add_argument('checkpoint', help='checkpoint file')
     parser.add_argument('input_dir', help='directory of the input video')
     parser.add_argument('output_dir', help='directory of the output video')
+    parser.add_argument(
+        '--filename_tmpl',
+        default='{:08d}.png',
+        help='template of the file names')
     parser.add_argument(
         '--window_size',
         type=int,
@@ -30,7 +34,7 @@ def main():
         args.config, args.checkpoint, device=torch.device('cuda', args.device))
 
     output = restoration_video_inference(model, args.input_dir,
-                                         args.window_size)
+                                         args.window_size, args.filename_tmpl)
     for i in range(0, output.size(1)):
         output_i = output[:, i, :, :, :]
         output_i = tensor2img(output_i)

mmedit/apis/restoration_video_inference.py

@@ -18,7 +18,7 @@ def pad_sequence(data, window_size):
     return data
 
 
-def restoration_video_inference(model, img_dir, window_size):
+def restoration_video_inference(model, img_dir, window_size, filename_tmpl):
     """Inference image with the model.
 
     Args:
@@ -31,11 +31,15 @@ def restoration_video_inference(model, img_dir, window_size):
     Returns:
         Tensor: The predicted restoration result.
     """
+
     device = next(model.parameters()).device  # model device
 
     # pipeline
     test_pipeline = [
-        dict(type='GenerateSegmentIndices', interval_list=[1]),
+        dict(
+            type='GenerateSegmentIndices',
+            interval_list=[1],
+            filename_tmpl=filename_tmpl),
         dict(
             type='LoadImageFromFileList',
             io_backend='disk',

mmedit/datasets/pipelines/augmentation.py

@@ -927,10 +927,12 @@ class GenerateSegmentIndices:
         interval_list (list[int]): Interval list for temporal augmentation.
             It will randomly pick an interval from interval_list and sample
             frame index with the interval.
+        filename_tmpl (str): Template for file name. Default: '{:08d}.png'.
     """
 
-    def __init__(self, interval_list):
+    def __init__(self, interval_list, filename_tmpl='{:08d}.png'):
         self.interval_list = interval_list
+        self.filename_tmpl = filename_tmpl
 
     def __call__(self, results):
         """Call function.
@@ -964,11 +966,11 @@ def __call__(self, results):
         lq_path_root = results['lq_path']
         gt_path_root = results['gt_path']
         lq_path = [
-            osp.join(lq_path_root, clip_name, f'{v:08d}.png')
+            osp.join(lq_path_root, clip_name, self.filename_tmpl.format(v))
             for v in neighbor_list
         ]
         gt_path = [
-            osp.join(gt_path_root, clip_name, f'{v:08d}.png')
+            osp.join(gt_path_root, clip_name, self.filename_tmpl.format(v))
             for v in neighbor_list
         ]