update tsv file construction guide

24623 · Jul 21, 2023 · 8cb9387 · 8cb9387
1 parent 4e87090
commit 8cb9387
Show file tree

Hide file tree

Showing 187 changed files with 1,056 additions and 0 deletions.
diff --git a/PREPRO.md b/PREPRO.md
@@ -55,3 +55,53 @@ wget -O body_pose_model.pth https://www.dropbox.com/sh/7xbup2qsn7vvjxo/AABaYNMvv
 python ./annotator/openpose/run.py --dataset_root /path/to/target/dataset
 ```
 
+
+
+
+
+### TSV preparation
+
+After running Grounded SAM and OpenPose, prepare the data structured as below:
+
+```
+Target Dataset
+└── 001/
+    ├── 0000.png
+    ├── 0001.png
+    ├── ...
+    ├────── groundsam
+    |       ├── 000001.png.mask.jpg
+    |       ├── 000002.png.mask.jpg
+    |       └── ...     
+    └────── openpose_json
+            ├── 000001.png.json
+            ├── 000002.png.json
+            └── ...
+└── 002/
+    ├── 0000.png
+    ├── 0001.png
+    ├── ...
+    ├────── groundsam
+    |       ├── 000001.png.mask.jpg
+    |       ├── 000002.png.mask.jpg
+    |       └── ...     
+    └────── openpose_json
+            ├── 000001.png.json
+            ├── 000002.png.json
+            └── ...
+```
+
+We also provide a preprocessed toy dataset as example. You may find the example dataset in folder `./tsv_example/toy_dataset`
+
+Run the following script to convert your dataset to tsv format
+
+```
+python ./tsv_example/create_custom_dataset_tsvs.py --split train --root_folder PATH_TO_YOUR_DATASET_FOLDER --output_folder PATH_TO_DESIRED_FOLDER 
+```
+
+For instance, `PATH_TO_YOUR_DATASET_FOLDER=./tsv_example/toy_dataset` and `PATH_TO_DESIRED_FOLDER=./tsv_example/toy_dataset/tsv`
+
+
+### TSV visualization
+
+Once the tsv files are generated, we provide a jupyter notebook for data visualization. Please refer to `visualize_tsv.ipynb` for details.
diff --git a/README.md b/README.md
@@ -14,6 +14,7 @@
 
 ## :fire: News
 
+* **[2023.07.21]** Update the [construction guide](https://github.com/Wangt-CN/DisCo/blob/main/PREPRO.md) of the TSV file.
 * **[2023.07.08]** Update the [Colab](https://colab.research.google.com/drive/1dErsSwSEdcAyP6V_mqayW0qdPuoHSz-j?usp=sharing) Demo (make sure our code/demo can be run on any machine)!
 * **[2023.07.03]** Provide the local demo deployment [example code](https://github.com/Wangt-CN/DisCo#-demo). Now you can try our demo on you own dev machine!
 * **[2023.07.03]** We update the Pre-training [tsv data](https://github.com/Wangt-CN/DisCo#1-human-attribute-pre-training).
@@ -145,6 +146,8 @@ Data Root
     ...
 ```
 
+*PS: If you want to use your own data resource but with our TSV data structure, please follow [PREPRO.MD](https://github.com/Wangt-CN/DisCo/blob/main/PREPRO.md) for reference. 
+
 <br><br/>
 
 
@@ -285,6 +288,8 @@ To run the visualization, just change `--do_train` to `--eval_visu` . You can al
 
 - Use Grounded-SAM and OpenPose to obtain human mask and human skeleton for each training image (See [PREPRO.MD](https://github.com/Wangt-CN/DisCo/blob/main/PREPRO.md) for more details)
 
+- For human-specific fine-tuning, we recommend to directly use raw image/mask/pose for training rather than build TSV file. If you still want to use TSV file structure to prepare your data,  please follow [PREPRO.MD](https://github.com/Wangt-CN/DisCo/blob/main/PREPRO.md) for reference.
+
 
 
 #### 2. Run the following script for human-specific fine-tuning:

diff --git a/tool/video/gen_gifs_for_fvd.py b/tool/video/gen_gifs_for_fvd.py
@@ -0,0 +1,114 @@
+import glob
+import math
+from collections import defaultdict
+from tqdm import tqdm
+from PIL import Image
+import numpy as np
+import imageio
+import threading
+import os
+import sys
+pythonpath = os.path.abspath(
+    os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
+print(pythonpath)
+sys.path.insert(0, pythonpath)
+
+class Preprocess(threading.Thread):
+    def __init__(self, data, begin_idx, end_idx, gif_root, frames_len, fps, use_tqdm=False, overwrite=False):
+        threading.Thread.__init__(self)
+        self.data = data
+        self.begin_idx = begin_idx
+        self.end_idx = end_idx
+        self.use_tqdm = use_tqdm
+        self.gif_root = gif_root
+        self.frames_len = frames_len
+        self.fps = fps
+        self.end_idx =  self.end_idx - self.frames_len
+        self.overwrite = overwrite
+
+    def run(self):
+        vid = "_".join(os.path.splitext(os.path.basename(self.data[0]))[0].split("_")[:2])
+        it = tqdm(range(self.begin_idx, self.end_idx), desc=f"{vid}") if self.use_tqdm else range(self.begin_idx, self.end_idx)
+        for idx in it:
+            file = self.data[idx]
+            image = imageio.imread(file)
+            base_name = os.path.splitext(os.path.basename(file))[0]
+            if os.path.exists(os.path.join(self.gif_root, f'{base_name}.gif')) and not self.overwrite:
+                continue
+            images = [image]
+            for i in range(1, self.frames_len):
+                try:
+                    file = self.data[idx+i]
+                    image = imageio.imread(file)
+                except Exception as e:
+                    print(e, f"curr_idx:{idx+i}", f"begin_idx:{self.begin_idx}", f"end_idx:{self.end_idx}", f"data_len: {len(self.data)}")
+                images.append(image)
+            imageio.mimsave(os.path.join(self.gif_root, f'{base_name}.gif'), images, fps=self.fps)
+
+
+def compose_gif(img_path, gif_frames, fps, num_workers, overwrite=False):
+    if img_path[-1] == "/":
+        img_path = img_path[:-1]
+    base = os.path.basename(img_path) + "_gif"
+    root = os.path.dirname(img_path)
+    gif_root = os.path.join(root, base)
+    if not os.path.exists(gif_root):
+        os.mkdir(gif_root)
+    inst_names = glob.glob(f"{img_path}/*.png")
+    inst_names = sorted(inst_names)
+
+    vid2files = defaultdict(list)
+    for path in inst_names:
+        vid = "_".join(os.path.splitext(os.path.basename(path))[0].split("_")[:2])
+        vid2files[vid].append(path)
+
+    for vid, inst_paths in vid2files.items():
+        pool = []
+        per_num = math.ceil(len(inst_paths) / num_workers)
+        cur_idx = 0
+        for idx in range(num_workers):
+            begin_idx = cur_idx
+            end_idx = cur_idx + per_num
+            if end_idx > len(inst_paths):
+                end_idx = len(inst_paths)
+            pool.append(
+                Preprocess(
+                    inst_paths, begin_idx, end_idx, gif_root,
+                    gif_frames, fps, idx==0, overwrite=overwrite))
+            cur_idx = end_idx
+
+        for idx in range(num_workers):
+            pool[idx].start()
+
+        for idx in range(num_workers):
+            pool[idx].join()
+
+    return gif_root
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    #parser.add_argument('--root_dir', type=str, default="/home/t-mni/blob/chenfeipremiumblock/data/G")
+    parser.add_argument('--root_dir', type=str, default="/f_ndata/G")
+    parser.add_argument('--path_gen', type=str, default="sdm/SDM_KLF8_S512_MSCOCO/eval_visu/pred")
+    parser.add_argument('--path_gt', type=str, default="dataset/mscoco/val2017")
+    parser.add_argument('--gif_frames', type=int, default=4)
+    parser.add_argument('--gif_fps', type=float, default=1.6)
+    parser.add_argument('--overwrite', type=bool, default=False)
+    parser.add_argument('--num_workers', type=int, default=16)
+
+    args = parser.parse_args()
+
+    if args.root_dir not in args.path_gen:
+        args.path_gen = os.path.join(args.root_dir, args.path_gen)
+    if args.root_dir not in args.path_gt:
+        args.path_gt = os.path.join(args.root_dir, args.path_gt)
+    print(f"compose_gif for pred")
+    path_gen_gif = compose_gif(args.path_gen, args.gif_frames,
+        args.gif_fps, args.num_workers, overwrite=args.overwrite)
+    print(f"compose_gif for gt")
+    path_gt_gif = compose_gif(args.path_gt, args.gif_frames,
+        args.gif_fps, args.num_workers, overwrite=args.overwrite)
+    print(path_gen_gif, path_gt_gif)
diff --git a/tool/video/gen_vid.py b/tool/video/gen_vid.py
@@ -0,0 +1,29 @@
+import cv2
+import numpy as np
+import os
+from tqdm import tqdm
+
+video_img_source_path = ''
+
+
+fps = 30
+size = (256, 256)
+codec = cv2.VideoWriter_fourcc(*"mp4v")
+
+# Initialize video writer
+video_writer = cv2.VideoWriter("{}_vid_output.mp4".format(video_img_source_path), codec, fps, size)
+folder_path = video_img_source_path
+# Get a list of all files in the folder
+files = os.listdir(folder_path)
+# Filter the list to only include image filesd
+image_files = [file for file in files if file.endswith(".jpg") or file.endswith(".png") or file.endswith(".jpeg")]
+num_frm = len(image_files)
+
+for image_fname_ in image_files:
+    image_fname = '{}/{}'.format(folder_path,image_fname_)
+    frame = cv2.imread(image_fname)
+    # Write frames to video
+    video_writer.write(frame)
+
+# Release resources
+video_writer.release()
diff --git a/tool/video/gen_vid_command.py b/tool/video/gen_vid_command.py
@@ -0,0 +1,34 @@
+import cv2
+import numpy as np
+import os
+import argparse
+from tqdm import tqdm
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Get Video')
+    parser.add_argument('--img_path', default=None, type=str, help='image frame path')
+    args = parser.parse_args()
+
+    video_img_source_path = args.img_path
+    fps = 30
+    size = (256, 256)
+    codec = cv2.VideoWriter_fourcc(*"mp4v")
+
+    # Initialize video writer
+    video_writer = cv2.VideoWriter("{}_vid_output.mp4".format(video_img_source_path), codec, fps, size)
+    folder_path = video_img_source_path
+    # Get a list of all files in the folder
+    files = os.listdir(folder_path)
+    # Filter the list to only include image filesd
+    image_files = [file for file in files if file.endswith(".jpg") or file.endswith(".png") or file.endswith(".jpeg")]
+    num_frm = len(image_files)
+
+    for image_fname_ in image_files:
+        image_fname = '{}/{}'.format(folder_path, image_fname_)
+        frame = cv2.imread(image_fname)
+        # Write frames to video
+        video_writer.write(frame)
+
+    # Release resources
+    video_writer.release()
diff --git a/tool/video/gen_vid_folder.py b/tool/video/gen_vid_folder.py
@@ -0,0 +1,55 @@
+import cv2
+import numpy as np
+import os
+from tqdm import tqdm
+# Define video settings
+
+
+folder_name = ''
+save_folder_name = ''
+
+sub_folder_name_list = os.listdir(folder_name)
+
+for idx, sub_folder_name in enumerate(sub_folder_name_list):
+    print(sub_folder_name)
+    # video_img_source_path = os.path.join(folder_name, sub_folder_name, 'eval_step_7080', 'pred_gs3.0_scale-cond1.0-ref1.0')
+    video_img_source_path = os.path.join(folder_name, sub_folder_name, 'eval_step_7200', 'pred_gs3.0_scale-cond1.0-ref1.0')
+    # video_img_source_path = os.path.join(folder_name, sub_folder_name, 'eval_step_17790', 'pred_gs3.0_scale-cond1.0-ref1.0')
+    if save_folder_name is None:
+        save_path = video_img_source_path
+    else: # use specific save folder
+        save_path = os.path.join(save_folder_name, sub_folder_name, 'pred_gs3.0_scale-cond1.0-ref1.0')
+        os.makedirs(save_path, exist_ok=True)
+
+    # if exists
+    video_save_path = "{}_vid_output.mp4".format(save_path)
+    if os.path.isfile(video_save_path) and os.path.getsize(video_save_path)>10000:
+        print('has done, skip this')
+        continue
+
+    fps = 30
+    size = (256, 256)
+    codec = cv2.VideoWriter_fourcc(*"mp4v")
+
+    # Get a list of all files in the folder
+    folder_path = video_img_source_path
+    try:
+        files = os.listdir(folder_path)
+    except:
+        print('load frame wrong, skip this')
+        continue
+
+    # Initialize video writer
+    video_writer = cv2.VideoWriter(video_save_path, codec, fps, size)
+    # Filter the list to only include image filesd
+    image_files = [file for file in files if file.endswith(".jpg") or file.endswith(".png") or file.endswith(".jpeg")]
+    num_frm = len(image_files)
+
+    for image_fname_ in image_files:
+        image_fname = '{}/{}'.format(folder_path,image_fname_)
+        frame = cv2.imread(image_fname)
+        # Write frames to video
+        video_writer.write(frame)
+
+    # Release resources
+    video_writer.release()
diff --git a/tool/video/reencode_vid_folder_forvis.py b/tool/video/reencode_vid_folder_forvis.py
@@ -0,0 +1,26 @@
+
+import os
+
+
+folder_name = ''
+save_folder_name = ''
+
+sub_folder_name_list = os.listdir(folder_name)
+sub_folder_name_list = [file for file in sub_folder_name_list if not 'json' in file]
+sub_folder_name_list.sort()
+
+for idx, sub_folder_name in enumerate(sub_folder_name_list):
+    print(sub_folder_name)
+
+    video_source_path = os.path.join(folder_name, sub_folder_name, 'cond_vid_output.mp4')
+    video_target_path = os.path.join(folder_name, sub_folder_name, 'cond_vid_output_reencode.mp4')
+
+    assert os.path.exists(video_source_path)
+
+    if os.path.exists(video_target_path):
+        print('remove the previous file')
+        os.remove(video_target_path)
+    os.system(f'ffmpeg -i {video_source_path} -c:v libx264 -c:a aac -strict -2 {video_target_path}')
+
+
+
diff --git a/tool/video/runsh/process_video.sh b/tool/video/runsh/process_video.sh
@@ -0,0 +1,18 @@
+export PYTHONUNBUFFERED=TRUE
+python tool/video/gen_vid_folder_forvis.py --start_idx 0 --end_idx 200 > tool/video/runsh/cond_log1.txt &
+python tool/video/gen_vid_folder_forvis.py --start_idx 200 --end_idx 400 > tool/video/runsh/cond_log2.txt &
+python tool/video/gen_vid_folder_forvis.py --start_idx 400 --end_idx 600 > tool/video/runsh/cond_log3.txt &
+python tool/video/gen_vid_folder_forvis.py --start_idx 600 --end_idx 800 > tool/video/runsh/cond_log4.txt &
+python tool/video/gen_vid_folder_forvis.py --start_idx 800 --end_idx 1000 > tool/video/runsh/cond_log5.txt &
+python tool/video/gen_vid_folder_forvis.py --start_idx 1000 --end_idx 1200 > tool/video/runsh/cond_log6.txt &
+python tool/video/gen_vid_folder_forvis.py --start_idx 1200 --end_idx 1400 > tool/video/runsh/cond_log7.txt &
+python tool/video/gen_vid_folder_forvis.py --start_idx 1400 --end_idx 1700 > tool/video/runsh/cond_log8.txt &
+
+#export PYTHONUNBUFFERED=TRUE
+#python tool/video/gen_vid_folder_forvis.py --start_idx 1400 --end_idx 1600 > tool/video/runsh/cond_log8.txt &
+#python tool/video/gen_vid_folder_forvis.py --start_idx 1600 --end_idx 1800 > tool/video/runsh/cond_log9.txt &
+#python tool/video/gen_vid_folder_forvis.py --start_idx 1800 --end_idx 2000 > tool/video/runsh/cond_log10.txt &
+#python tool/video/gen_vid_folder_forvis.py --start_idx 2000 --end_idx 2200 > tool/video/runsh/cond_log11.txt &
+#python tool/video/gen_vid_folder_forvis.py --start_idx 2200 --end_idx 2400 > tool/video/runsh/cond_log12.txt &
+#python tool/video/gen_vid_folder_forvis.py --start_idx 2400 --end_idx 2600 > tool/video/runsh/cond_log13.txt &
+#python tool/video/gen_vid_folder_forvis.py --start_idx 2600 --end_idx 2800 > tool/video/runsh/cond_log14.txt &
diff --git a/tool/video/runsh/process_video_demo.sh b/tool/video/runsh/process_video_demo.sh
@@ -0,0 +1,13 @@
+export PYTHONUNBUFFERED=TRUE
+python tool/video/gen_vid_folder_forvis.py --start_idx 0 --end_idx 200 --gen_idx 1 > tool/video/runsh/cond_log11.txt &
+python tool/video/gen_vid_folder_forvis.py --start_idx 200 --end_idx 400 --gen_idx 1 > tool/video/runsh/cond_log12.txt &
+python tool/video/gen_vid_folder_forvis.py --start_idx 400 --end_idx 600 --gen_idx 1 > tool/video/runsh/cond_log13.txt &
+#python tool/video/gen_vid_folder_forvis.py --start_idx 600 --end_idx 800 > tool/video/runsh/cond_log4.txt &
+#python tool/video/gen_vid_folder_forvis.py --start_idx 800 --end_idx 1000 > tool/video/runsh/cond_log5.txt &
+#python tool/video/gen_vid_folder_forvis.py --start_idx 1000 --end_idx 1200 > tool/video/runsh/cond_log6.txt &
+#python tool/video/gen_vid_folder_forvis.py --start_idx 1200 --end_idx 1400 > tool/video/runsh/cond_log7.txt &
+#python tool/video/gen_vid_folder_forvis.py --start_idx 1400 --end_idx 1700 > tool/video/runsh/cond_log8.txt &
+
+python tool/video/gen_vid_folder_forvis.py --start_idx 0 --end_idx 200 --gen_idx 0 > tool/video/runsh/log11.txt &
+python tool/video/gen_vid_folder_forvis.py --start_idx 200 --end_idx 400 --gen_idx 0 > tool/video/runsh/log12.txt &
+python tool/video/gen_vid_folder_forvis.py --start_idx 400 --end_idx 600 --gen_idx 0 > tool/video/runsh/log13.txt &