def main()

in scripts/data_preparation/prepare_cityscapes.py [0:0]

• 59 lines of code
• 11 McCabe index (conditional complexity)

def main(args):
    print("Loading Cityscapes from", args.root_dir)
    num_stuff, num_thing = _get_meta()

    # Prepare directories
    img_dir = path.join(args.out_dir, "img")
    _ensure_dir(img_dir)
    msk_dir = path.join(args.out_dir, "msk")
    _ensure_dir(msk_dir)
    lst_dir = path.join(args.out_dir, "lst")
    _ensure_dir(lst_dir)
    coco_dir = path.join(args.out_dir, "coco")
    _ensure_dir(coco_dir)

    # COCO-style category list
    coco_categories = []
    for lbl in cs_labels:
        if lbl.trainId != 255 and lbl.trainId != -1 and lbl.hasInstances:
            coco_categories.append({
                "id": lbl.trainId,
                "name": lbl.name
            })

    # Process splits
    images = []
    for split, (split_img_subdir, split_msk_subdir) in _SPLITS.items():
        print("Converting", split, "...")

        img_base_dir = path.join(args.root_dir, split_img_subdir)
        msk_base_dir = path.join(args.root_dir, split_msk_subdir)
        img_list = _get_images(msk_base_dir)

        # Write the list file
        with open(path.join(lst_dir, split + ".txt"), "w") as fid:
            fid.writelines(img_id + "\n" for _, img_id, _ in img_list)

        # Convert to COCO detection format
        coco_out = {
            "info": {"version": "1.0"},
            "images": [],
            "categories": coco_categories,
            "annotations": []
        }

        # Process images in parallel
        worker = _Worker(img_base_dir, msk_base_dir, img_dir, msk_dir)
        with Pool(initializer=_init_counter, initargs=(_Counter(0),)) as pool:
            total = len(img_list)
            for img_meta, coco_img, coco_ann in tqdm.tqdm(pool.imap(worker, img_list, 8), total=total):
                images.append(img_meta)

                # COCO annotation
                coco_out["images"].append(coco_img)
                coco_out["annotations"] += coco_ann

        # Write COCO detection format annotation
        with open(path.join(coco_dir, split + ".json"), "w") as fid:
            json.dump(coco_out, fid)

    # Write meta-data
    print("Writing meta-data")
    meta = {
        "images": images,
        "meta": {
            "num_stuff": num_stuff,
            "num_thing": num_thing,
            "categories": [],
            "palette": [],
            "original_ids": []
        }
    }

    for lbl in cs_labels:
        if lbl.trainId != 255 and lbl.trainId != -1:
            meta["meta"]["categories"].append(lbl.name)
            meta["meta"]["palette"].append(lbl.color)
            meta["meta"]["original_ids"].append(lbl.id)

    with open(path.join(args.out_dir, "metadata.bin"), "wb") as fid:
        umsgpack.dump(meta, fid, encoding="utf-8")