def main()

in siammot/data/ingestion/ingest_mot.py [0:0]
62 lines of code
12 McCabe index (conditional complexity)

def main(args, description="Initial ingestion", det_options=None, mot17=True):
    if mot17:
        if det_options is not None and not all(x in DET_OPTIONS for x in det_options):
            raise ValueError("Det options were {} but must be only: {}".format(det_options, DET_OPTIONS))
        if det_options is None:
            det_options = DET_OPTIONS
    else:
        print("Ingesting MOT15, ignoring det options {}".format(det_options))
        det_options = [""]

    dataset_path = args.dataset_path
    out_filename = args.anno_name

    out_dataset = GluonCVMotionDataset(out_filename, dataset_path, load_anno=False)
    metadata = {
        FieldNames.DESCRIPTION: description,
        FieldNames.DATE_MODIFIED: str(datetime.datetime.now()),
    }
    out_dataset.metadata = metadata

    splits = {
        "train": os.path.join(out_dataset.data_root_path, "train"),
        "test": os.path.join(out_dataset.data_root_path, "test"), # No gt for MOT test
    }

    for det_option in det_options:
        for split_name, split_path in splits.items():
            subdirs = glob.glob(os.path.join(split_path, "*" + det_option))
            for i, subdir in enumerate(subdirs):
                vid_id = os.path.basename(subdir)
                vid_path = os.path.join(split_path, subdir)

                sample = DataSample(vid_id)

                if mot17:
                    info_path = os.path.join(vid_path, "seqinfo.ini")
                    config = configparser.ConfigParser()
                    config.read(info_path)
                    seq_conf = config["Sequence"]
                    fps = float(seq_conf['frameRate'])
                    num_frames = int(seq_conf['seqLength'])
                    width = int(seq_conf['imWidth'])
                    height = int(seq_conf['imHeight'])
                else:
                    # Assume 30 fps
                    fps = 30
                    im_paths = glob.glob(os.path.join(vid_path, "img1", "*.jpg"))
                    num_frames = len(im_paths)
                    im_example = Image.open(im_paths[0])
                    width = im_example.width
                    height = im_example.height

                rel_base_dir = vid_path.replace(out_dataset.data_root_path, "").lstrip(os.path.sep)
                rel_base_dir = os.path.join(rel_base_dir, "img1")
                metadata = {
                    FieldNames.DATA_PATH: rel_base_dir,
                    FieldNames.FPS: fps,
                    FieldNames.NUM_FRAMES: num_frames,
                    FieldNames.RESOLUTION: {"width": width, "height": height},
                }
                sample.metadata = metadata

                gt_path = os.path.join(vid_path, "gt/gt.txt")
                det_path = os.path.join(vid_path, "det/det.txt")
                has_gt = os.path.exists(gt_path)
                anno_path = gt_path if has_gt else det_path

                sample = sample_from_mot_csv(anno_path, fps, sample, mot17, has_gt)

                out_dataset.add_sample(sample)

                print("Done {} sample {}/{}, {}".format(split_name, i+1, len(subdirs), vid_id))

    out_dataset.dump()

    return out_dataset