From 3aacd8d5a75d2a29d8239b72673ccfde3aca8eba Mon Sep 17 00:00:00 2001 From: julie wang Date: Tue, 3 Sep 2024 09:11:50 +0200 Subject: [PATCH] feat(ml):add mini dataset for unite test of video series dataloader --- ...ed_temporal_labeled_mask_online_dataset.py | 84 +++++++++++++------ scripts/run_tests.sh | 10 +++ tests/test_run_vid_diffusion_online.py | 8 +- 3 files changed, 75 insertions(+), 27 deletions(-) diff --git a/data/self_supervised_temporal_labeled_mask_online_dataset.py b/data/self_supervised_temporal_labeled_mask_online_dataset.py index 896515b88..c56c385be 100644 --- a/data/self_supervised_temporal_labeled_mask_online_dataset.py +++ b/data/self_supervised_temporal_labeled_mask_online_dataset.py @@ -1,9 +1,9 @@ import os import random import re - +import bisect import torch - +from collections import OrderedDict from data.base_dataset import BaseDataset, get_transform_list from data.image_folder import make_labeled_path_dataset from data.online_creation import crop_image @@ -61,6 +61,32 @@ def __init__(self, opt, phase, name=""): self.A_size = len(self.A_img_paths) # get the size of dataset A + # dataset form img(bbox)/vid_series/vid_series_#frame.png(.txt) + # a ordered list with all video series paths + self.vid_series_paths = list( + OrderedDict.fromkeys(os.path.dirname(path) for path in self.A_img_paths) + ) + # Initialize a dictionary to count how many available paths belong to each directory + self.frames_counts = { + vid_serie: -self.num_frames * self.frame_step + for vid_serie in self.vid_series_paths + } + # Loop through self.A_img_paths and count the occurrences of each directory + for path in self.A_img_paths: + dirname = os.path.dirname(path) + if dirname in self.vid_series_paths: + self.frames_counts[dirname] += 1 + + # Store cumulative sums of available frames in the order of video series. + self.cumulative_sums = [] + cumulative_sum = 0 + # Iterate through each video series in frames_counts to compute the cumulative sum of available frame. + for vid_serie in self.vid_series_paths: + count_minus = self.frames_counts[vid_serie] + if count_minus > 0: + cumulative_sum += count_minus + self.cumulative_sums.append(cumulative_sum) + def get_img( self, A_img_path, @@ -72,30 +98,38 @@ def get_img( index=None, ): # all params are unused - while True: - + if len(self.frames_counts) == 1: # single video mario index_A = random.randint(0, self.range_A - 1) - - images_A = [] - labels_A = [] - - ref_A_img_path = self.A_img_paths[index_A] - ref_name_A = ref_A_img_path.split("/")[-1][: self.num_common_char] - ref_A_name = ref_A_img_path.split("/")[-1] # fullname of the ref_A - - vid_series_path = os.path.dirname(ref_A_img_path) - vid_series = vid_series_path.split("/")[-1] - if ref_A_name.startswith( - vid_series - ): # dataset contains different video in form of img/vid_series/vid_seriesframe.jpg - series_count = sum(vid_series_path in path for path in self.A_img_paths) - frame_num = int(ref_A_name[len(vid_series) : -4]) # remove ".jpg" - if frame_num < (series_count - self.num_frames): - break - else: - print("Condition not met, generating a new index_A...") - else: # dataset from one video in form of img/frames.jpg - break + else: # video series + range_A = self.cumulative_sums[ + -1 + ] # total number of frames that can be used as a starting frame + random_A = random.randint( + 0, range_A - 1 + ) # chose one frame from available video series + + # according to the selected_index, get the video series and frame number + selected_index = bisect.bisect_left(self.cumulative_sums, random_A) + selected_vid = self.vid_series_paths[selected_index] + if selected_index > 0: + frame_num = random_A - self.cumulative_sums[selected_index - 1] + else: + frame_num = random_A + + filtered_paths = [ + path + for path in self.A_img_paths + if os.path.dirname(path) == selected_vid + ] + + # Get path and index_A + selected_path = filtered_paths[frame_num - 1] + index_A = self.A_img_paths.index(selected_path) + + images_A = [] + labels_A = [] + ref_A_img_path = self.A_img_paths[index_A] + ref_name_A = ref_A_img_path.split("/")[-1][: self.num_common_char] for i in range(self.num_frames): cur_index_A = index_A + i * self.frame_step diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh index ffd584197..ee9d17014 100755 --- a/scripts/run_tests.sh +++ b/scripts/run_tests.sh @@ -205,6 +205,16 @@ fi ####### vid diffusion tests echo "Running vid diffusion training tests" +URL=https://www.joligen.com/datasets/bdd100K_vid.zip +ZIP_FILE=$DIR/bdd100K_vid.zip +TARGET_MASK_SEM_ONLINE_DIR=$DIR/bdd100k_vid +wget -N $URL -O $ZIP_FILE +mkdir $TARGET_MASK_SEM_ONLINE_DIR +unzip $ZIP_FILE -d $DIR +rm $ZIP_FILE +ln -s $TARGET_MASK_SEM_ONLINE_DIR/trainA $TARGET_MASK_SEM_ONLINE_DIR/testA +ln -s $TARGET_MASK_SEM_ONLINE_DIR/trainB $TARGET_MASK_SEM_ONLINE_DIR/testB + python3 -m pytest --rootdir ${current_dir} -p no:cacheprovider -s "${current_dir}/../tests/test_run_vid_diffusion_online.py" --dataroot "$TARGET_MASK_SEM_ONLINE_DIR" OUT=$? diff --git a/tests/test_run_vid_diffusion_online.py b/tests/test_run_vid_diffusion_online.py index 18da71c51..35c562c72 100644 --- a/tests/test_run_vid_diffusion_online.py +++ b/tests/test_run_vid_diffusion_online.py @@ -62,9 +62,13 @@ def test_vid_diffusion_online(dataroot): json_like_dict["dataroot"] = dataroot json_like_dict["checkpoints_dir"] = "/".join(dataroot.split("/")[:-1]) - with open( - os.path.join(json_like_dict["dataroot"], "trainA", "paths.txt") + os.path.join( + json_like_dict["checkpoints_dir"], + dataroot.split("/")[-1], + "trainA", + "paths.txt", + ) ) as paths_file: line = paths_file.readline().strip().split(" ") infer_options["img_in"] = os.path.join(json_like_dict["dataroot"], line[0])