def read_gt()

in videoalignment/datasets.py [0:0]

• 121 lines of code
• 25 McCabe index (conditional complexity)

def read_gt(dataset):
    check_data(dataset)
    root_dir = DATASETS[dataset]
    if dataset in ["Climbing", "Madonna"]:
        videos = []
        gt_file = os.path.join(root_dir, f"gt_{dataset.lower()}.align")
        for line in open(gt_file, "r"):
            chunks = line.split(" ")
            video = chunks[0]
            if dataset == "Madonna":
                video = video.split("/")[1].split(".")[0]
            else:
                video = video.split(".")[0]
            videos.append(
                {
                    "video": video,
                    "begin": float(chunks[1]),
                    "end": float(chunks[2]),
                    "global": float(chunks[3]),
                }
            )
        pairs = []
        for video_a in videos:
            video_o = videos[videos.index(video_a) + 1 :]
            for video_b in video_o:
                extent1 = np.around(
                    np.arange(
                        video_a["global"] + video_a["begin"],
                        video_a["global"] + video_a["end"],
                        0.01,
                    ),
                    2,
                )
                extent2 = np.around(
                    np.arange(
                        video_b["global"] + video_b["begin"],
                        video_b["global"] + video_b["end"],
                        0.01,
                    ),
                    2,
                )
                offset = compute_offset(video_a, video_b)
                if np.intersect1d(extent1, extent2).size > 0:
                    pairs.append({"videos": [video_a, video_b], "offset": offset})

    elif dataset == "VCDB":
        annotation_dir = os.path.join(root_dir, "annotation")
        videos = []
        pairs = []
        for filename in os.listdir(annotation_dir):
            video_file = filename.split(".")[0]
            for line_i, line in enumerate(open("%s/%s" % (annotation_dir, filename))):

                chunks = line.split(",")
                begin_a = chunks[2].split(":")[::-1]
                begin_a = sum(int(t) * 60 ** i for (i, t) in enumerate(begin_a))
                end_a = chunks[3].split(":")[::-1]
                end_a = sum(int(t) * 60 ** i for (i, t) in enumerate(end_a))
                begin_b = chunks[4].split(":")[::-1]
                begin_b = sum(int(t) * 60 ** i for (i, t) in enumerate(begin_b))
                end_b = chunks[5].split(":")[::-1]
                end_b = sum(int(t) * 60 ** i for (i, t) in enumerate(end_b))

                video_a = {
                    "video": "%s/%s" % (video_file, chunks[0].split(".")[0]),
                    "global": -begin_a,
                    "begin": begin_a,
                    "end": end_a,
                }
                video_b = {
                    "video": "%s/%s" % (video_file, chunks[1].split(".")[0]),
                    "global": -begin_b,
                    "begin": begin_b,
                    "end": end_b,
                }

                offset = compute_offset(video_a, video_b)

                if video_a not in videos:
                    videos.append(video_a)
                else:
                    video_a = videos[videos.index(video_a)]
                if video_b not in videos:
                    videos.append(video_b)
                else:
                    video_b = videos[videos.index(video_b)]

                pairs.append({"videos": [video_a, video_b], "offset": offset})
    elif dataset == "EVVE":
        annotation_dir = os.path.join(root_dir, "annotations")
        videos = []
        pairs = []
        durations = os.path.join(annotation_dir, "durations.pkl")
        durations = pkl.load(open(durations, "rb"))
        for fname in os.listdir(annotation_dir):
            if not fname.endswith(".dat"):
                continue
            event = fname.split(".")[0]
            this_videos = []
            for l_i, l in enumerate(open(os.path.join(annotation_dir, fname), "r")):
                vidname, gt, split = l.split()
                vidname = "%s/%s" % (event, vidname)
                duration = durations[vidname]
                gt = int(gt)
                video = {
                    "video": vidname,
                    "begin": 0,
                    "end": duration,
                    "global": 0,
                    "split": split,
                    "pos_null": gt,
                    "event": event,
                }
                this_videos.append(video)

            videos.extend(this_videos)

            for video_a in this_videos:
                video_o = this_videos[this_videos.index(video_a) + 1 :]
                for video_b in video_o:
                    if (
                        video_a["split"] == video_b["split"]
                        and video_a["pos_null"] == 1
                        and video_b["pos_null"] == 1
                    ):
                        pairs.append({"videos": [video_a, video_b], "offset": -1})
    else:
        raise NotImplementedError
    return videos, pairs