def split_train_val()

in videoalignment/datasets.py [0:0]

• 22 lines of code
• 14 McCabe index (conditional complexity)

    def split_train_val(self):
        if self.phase == "all":
            self.videos = self.gt_all_videos
            self.overlapping_pairs = self.gt_all_overlapping_pairs
        elif self.phase == "train":
            videos_train = [v for v in self.gt_all_videos if v["split"] == "database"]
            pairs_train = [
                p
                for p in self.gt_all_overlapping_pairs
                if p["videos"][0] in videos_train and p["videos"][1] in videos_train
            ]
            self.videos = videos_train
            self.overlapping_pairs = pairs_train
        elif self.phase == "val":
            videos_val = [v for v in self.gt_all_videos if v["split"] == "query"]

            pairs_val = [
                p
                for p in self.gt_all_overlapping_pairs
                if p["videos"][0] in videos_val and p["videos"][1] in videos_val
            ]
            self.videos = videos_val
            self.overlapping_pairs = pairs_val