dataset.py

import typer
import fiftyone as fo
from fiftyone import ViewField as F
from pathlib import Path
from pycocotools.coco import COCO
from loguru import logger
import cv2
import shutil
import os
import random
from collections import defaultdict
import csv


DEFAULT_EXCLUDE_NAME = "Ellen"
DEFAULT_INS_TRAIN = "instances_Train.json"
DEFAULT_INS_TEST = "instances_Test.json"

app = typer.Typer()


@app.command()
def newsplit(
    in_dir: str,
    train_json=DEFAULT_INS_TRAIN,
    test_json=DEFAULT_INS_TEST,
    exclude_name=DEFAULT_EXCLUDE_NAME,
):
    """
    Merge the train and test datasets,
    and then split them into new train/test by leaving one person out.
    """

    # load the dataset
    logger.info("Loading datasets...")
    ds1 = fo.Dataset.from_dir(
        dataset_type=fo.types.COCODetectionDataset,
        data_path=Path(in_dir) / "images",
        labels_path=Path(in_dir) / "annotations" / train_json,
    )
    ds2 = fo.Dataset.from_dir(
        dataset_type=fo.types.COCODetectionDataset,
        data_path=Path(in_dir) / "images",
        labels_path=Path(in_dir) / "annotations" / test_json,
    )

    logger.info(f"[Before] Num samples in train: {len(ds1)}")
    logger.info(f"[Before] Num samples in test: {len(ds2)}")

    # merge the datasets
    ds1.merge_samples(ds2)

    # generate the new split
    logger.info(f"Excluding name in filepath as train set: {exclude_name}")
    new_train_view = ds1.match(~F("filepath").re_match(exclude_name))
    new_test_view = ds1.match(F("filepath").re_match(exclude_name))
    assert len(new_train_view) + len(new_test_view) == len(ds1)
    logger.info(f"[After] Num samples in train: {len(new_train_view)}")
    logger.info(f"[After] Num samples in test: {len(new_test_view)}")
    train_counts = new_train_view.count_values("detections.detections.label")
    test_counts = new_test_view.count_values("detections.detections.label")
    logger.info(f"[After] Train counts: {train_counts}")
    logger.info(f"[After] Test counts: {test_counts}")

    # export the new split
    logger.info("Exporting new train/test...")
    new_train_p = Path(in_dir) / "annotations" / f"new_train_no-{exclude_name}.json"
    new_test_p = Path(in_dir) / "annotations" / f"new_test_{exclude_name}.json"
    new_train_view.export(
        dataset_type=fo.types.COCODetectionDataset,
        labels_path=new_train_p,
        label_field="segmentations",
        classes=ds1.default_classes,
        abs_paths=True,
    )
    new_test_view.export(
        dataset_type=fo.types.COCODetectionDataset,
        labels_path=new_test_p,
        label_field="segmentations",
        classes=ds2.default_classes,
        abs_paths=True,
    )
    logger.info(f"Exported new train: {new_train_p}")
    logger.info(f"Exported new test: {new_test_p}")


def _normalize(img_size, xy_s):
    assert len(xy_s) % 2 == 0
    normalized_xy_s = []
    dw = 1.0 / (img_size[0])
    dh = 1.0 / (img_size[1])
    for i in range(len(xy_s)):
        p = xy_s[i]
        p = p * dw if i % 2 == 0 else p * dh
        assert p <= 1.0 and p >= 0.0, f"{p} should < 1 and > 0"
        normalized_xy_s.append(p)
    return normalized_xy_s


def _coco2yolo(coco_img_dir, coco_json_path, out_dir, bbox_only=False, rois=None):
    logger.info(f"Reading {Path(coco_json_path).name}...")
    coco = COCO(coco_json_path)

    cats = coco.loadCats(coco.getCatIds())
    cats = sorted(cats, key=lambda x: x["id"], reverse=False)
    assert cats[0]["id"] == 1, f"Assume cat id starts from 1, but got {cats[0]['id']}"
    logger.info(f"{len(cats)} categories: {[cat['name'] for cat in cats]}")

    img_ids = coco.getImgIds()
    prefix = Path(coco_json_path).stem.split("_")[-1].lower()  # either train or test

    # create output directories
    target_txt_r = Path(out_dir) / prefix / "labels"
    target_img_r = Path(out_dir) / prefix / "images"
    target_txt_r.mkdir(parents=True, exist_ok=False)
    target_img_r.mkdir(parents=True, exist_ok=False)

    logger.info(f"Num of imgs: {len(img_ids)}")

    n_imgs_no_annos = 0
    num_zero_area = 0
    for img_id in img_ids:
        img = coco.loadImgs(img_id)[0]
        img_p = Path(coco_img_dir) / img["file_name"]
        assert img_p.exists(), f"{img_p} does not exist"

        anno_ids = coco.getAnnIds(imgIds=img["id"])
        annos = coco.loadAnns(anno_ids)

        new_filename = f"{img['id']}_{img_p.stem}"

        out_img_p = target_img_r / (new_filename + img_p.suffix)

        # get roi for the image if any
        im_cv = cv2.imread(img_p.as_posix())
        im_width, im_height = im_cv.shape[1], im_cv.shape[0]
        roi = rois[(im_width, im_height)] if rois is not None else None
        has_roi = (rois is not None) and (roi is not None) and len(roi) == 4
        if not has_roi:
            # copy image to target dir
            shutil.copy(img_p, out_img_p)
        else:
            # crop the image to target dir
            assert len(roi) == 4, f"ROI should have 4 values, but got {roi}"
            cropped_img = im_cv[roi[1] : roi[1] + roi[3], roi[0] : roi[0] + roi[2]]
            cv2.imwrite(out_img_p.as_posix(), cropped_img)

        # bg imgs: only need to copy img, no need to create label file
        if len(annos) == 0:
            n_imgs_no_annos += 1
            continue

        # create the label txt file
        txt_p = Path(target_txt_r) / (new_filename + ".txt")
        if txt_p.exists():
            logger.warning(f"{txt_p} already exists, {img_p} skipped")
        txt_f = open(txt_p, "w")
        img = cv2.imread(img_p.as_posix())
        h, w, _ = img.shape

        # generate txt file for each image
        for ann in annos:
            cls_id = ann["category_id"] - 1  # yolov5 uses zero-based class idx

            # region bbox, for object detection
            if bbox_only:
                bbox = ann["bbox"]
                # convert coco to yolo: top-x, top-y, w, h -> center-x, center-y, w, h
                bbox_yolo = [
                    bbox[0] + bbox[2] / 2,
                    bbox[1] + bbox[3] / 2,
                    bbox[2],
                    bbox[3],
                ]
                n_bbox_p = " ".join([str(a) for a in _normalize((w, h), bbox_yolo)])
                txt_f.write(f"{cls_id} {n_bbox_p}{os.linesep}")
                continue
            # endregion

            # region seg, for instance segmentation
            seg = ann["segmentation"]
            if len(seg) > 1:
                # TODO: Investigate why sometimes there are multiple segs
                logger.warning(f"Skip {img_p} with {len(seg)} segs of {ann}")
                continue

            if len(seg) == 1:
                xy_s = seg[0]
                # handle roi if any
                if has_roi:
                    xy_s = [xy - roi[i % 2] for i, xy in enumerate(xy_s)]
                    w, h = roi[2], roi[3]
                    # remove the points outside of roi
                    new_xy_s = []
                    for i in range(0, len(xy_s), 2):
                        x, y = xy_s[i], xy_s[i + 1]
                        if x >= 0 and x <= w and y >= 0 and y <= h:
                            new_xy_s.extend([x, y])
                    xy_s = new_xy_s
                n_xy_s = _normalize((w, h), xy_s)
                seg_p = " ".join([str(a) for a in n_xy_s])
                txt_f.write(f"{cls_id} {seg_p}{os.linesep}")
            # endregion

            # region keypoint, for pose estimation
            if "keypoints" in ann:
                # skip area 0 keypoints which could cause yolov8 training error
                if int(ann["area"]) == 0:
                    num_zero_area += 1
                    continue
                kps = ann["keypoints"]
                bbox = ann["bbox"]
                # convert coco to yolo: top-x, top-y, w, h -> center-x, center-y, w, h
                bbox_yolo = [
                    bbox[0] + bbox[2] / 2,
                    bbox[1] + bbox[3] / 2,
                    bbox[2],
                    bbox[3],
                ]
                n_bbox_p = " ".join([str(a) for a in _normalize((w, h), bbox_yolo)])
                # normalize x,y of each keypoint and keep visibility as is
                n_kp = []
                for i in range(0, len(kps), 3):
                    n_kp.append(kps[i] / w)
                    n_kp.append(kps[i + 1] / h)
                    n_kp.append(kps[i + 2])
                n_kp_p = " ".join([str(a) for a in n_kp])
                txt_f.write(f"{cls_id} {n_bbox_p} {n_kp_p}{os.linesep}")
            # endregion
        txt_f.close()
        # remove empty label file which has no annos
        if txt_p.stat().st_size == 0:
            txt_p.unlink()
            n_imgs_no_annos += 1
    empty_ratio = 100 * float(n_imgs_no_annos) / len(img_ids)
    n_imgs_anns = len(img_ids) - n_imgs_no_annos
    logger.info(f"# imgs w anns: {n_imgs_anns} {(100-empty_ratio):.2f}%")
    logger.info(f"# imgs w/o anns: {n_imgs_no_annos} {empty_ratio:.2f}%")
    logger.info(f"# zero area kps: {num_zero_area}")
    txts = [f for f in target_txt_r.iterdir() if f.is_file()]
    imgs = [f for f in target_img_r.iterdir() if f.is_file()]
    assert (len(txts) + n_imgs_no_annos) == len(imgs) == len(img_ids)
    return target_img_r


@app.command(help="Convert COCO dataset to YOLOv5 format")
def coco2yolov5(
    in_dir: str,
    out_dir: str,
    split_val_ratio: float = 0.2,
    seed: int = 42,
    bbox_only: bool = False,
    crop_roi_file: str = None,
):
    """
    Convert COCO dataset to YOLOv5 format.
    Support 3 task types: object detection, instance segmentation, pose estimation.

    YOLOv5 seg labels are the same as detection labels, using txt files with one object per line.
    The difference is that instead of "class, xywh" they are "class xy1, xy2, xy3,...".
    Ref: https://github.com/ultralytics/yolov5/issues/10161#issuecomment-1315672357

    YOLOv5 keypoint labels is using txt files with one object per line.
        class cx cy w h x1 y1 v1 ... xn yn vn
    All coordinates are normalized by image width and height.
    vn (visibility): 0, 1, or 2 => not labeled, labeled but invisible, labeled and visible
    Ref: https://github.com/ultralytics/ultralytics/issues/1870#issuecomment-1498909244
    Example: https://ultralytics.com/assets/coco8-pose.zip
    """
    if Path(out_dir).exists():
        delete = typer.confirm(f"{out_dir} alread exists. Are you sure to delete it?")
        if not delete:
            logger.info("Not deleting")
            raise typer.Abort()
        shutil.rmtree(out_dir)
        logger.info(f"Deleted {Path(out_dir).name}")

    ann_dir_p = Path(in_dir) / "annotations"
    img_dir_p = Path(in_dir) / "images"
    assert ann_dir_p.exists(), f"{ann_dir_p} does not exist"
    assert img_dir_p.exists(), f"{img_dir_p} does not exist"

    # try to find the json files of train & test in annotations dir
    train_json_p = None
    test_json_p = None
    for f in ann_dir_p.iterdir():
        if f.stem.lower().endswith("train"):
            train_json_p = f
            logger.info(f"Found train json: {f.name}")
        elif f.stem.lower().endswith("test"):
            test_json_p = f
            logger.info(f"Found test json: {f.name}")
    # must have train, while test is optional
    assert train_json_p is not None, f"Cannot find train json in {ann_dir_p}"
    do_split = False
    if test_json_p is None:
        logger.warning("Cannot find test json in [in_dir]/annotations")
        do_split = typer.confirm("Do you want to split val from train?")

    # region handle ROIs
    rois = None
    if crop_roi_file is not None:
        roi_csv_p = Path(crop_roi_file)
        assert roi_csv_p.exists(), f"{roi_csv_p} does not exist"
        # read ROIs from csv, each image size should have one ROI
        rois = defaultdict(lambda: [], {})
        with open(roi_csv_p, "r") as f:
            for roi in csv.DictReader(f):
                ori_width = int(roi["ori_width"])
                ori_height = int(roi["ori_height"])
                roi_x = int(roi["roi_x"])
                roi_y = int(roi["roi_y"])
                roi_width = int(roi["roi_width"])
                roi_height = int(roi["roi_height"])

                key = (ori_width, ori_height)
                assert key not in rois, f"Duplicate ROI for {key}"
                rois[key] = [roi_x, roi_y, roi_width, roi_height]
    # endregion

    yolo_train_img_dir = None
    yolo_test_img_dir = None
    yolo_train_img_dir = _coco2yolo(img_dir_p, train_json_p, out_dir, bbox_only, rois)
    if test_json_p is not None:
        yolo_test_img_dir = _coco2yolo(img_dir_p, test_json_p, out_dir, bbox_only, rois)

    if do_split:
        yolo_test_img_dir = Path(out_dir) / "val" / "images"
        # randomly select 20% of train images
        train_imgs = [f for f in yolo_train_img_dir.iterdir() if f.is_file()]
        n_test = int(len(train_imgs) * split_val_ratio)
        logger.info(f"Split ratio {split_val_ratio}: {n_test} test images from train")
        # set random seed to make sure the same images are selected
        random.seed(seed)
        test_imgs = random.sample(train_imgs, n_test)
        # move test images to val/images
        yolo_test_img_dir.mkdir(parents=True, exist_ok=True)
        for f in test_imgs:
            shutil.move(str(f), str(yolo_test_img_dir))
        # move labels of test images to val/labels
        yolo_test_label_dir = Path(out_dir) / "val" / "labels"
        yolo_test_label_dir.mkdir(parents=True, exist_ok=True)
        for f in test_imgs:
            label_f = yolo_train_img_dir.parent / "labels" / f"{f.stem}.txt"
            if label_f.exists():
                shutil.move(str(label_f), str(yolo_test_label_dir))

    # region create yaml file

    logger.info(f"Reading {Path(train_json_p).name}...")
    train_coco = COCO(train_json_p)
    train_cats = train_coco.loadCats(train_coco.getCatIds())
    num_kps = [
        len(c["keypoints"])
        for c in train_cats
        if "keypoints" in c and len(c["keypoints"]) > 0
    ]
    # check if all categories have the same number of keypoints
    if len(num_kps) > 0:
        assert len(set(num_kps)) == 1, "Categories have different number of keypoints"
    logger.info(f"Number of keypoints: {set(num_kps)}")
    train_cats = [c["name"] for c in train_cats]
    # ensure having the same categories in the json of train & test
    # test_coco = COCO(test_json_p)
    # test_cats = test_coco.loadCats(test_coco.getCatIds())
    # test_cats = sorted(test_cats, key=lambda x: x["id"], reverse=False)
    # test_cats = [c["name"] for c in test_cats]
    # assert ",".join(train_cats) == ",".join(test_cats), "Categories mismatch"

    out_config_file = Path(out_dir) / "data.yaml"
    with open(out_config_file, "w") as f:
        if len(num_kps) > 0:
            f.write(f"kpt_shape: [{num_kps[0]},3]" + os.linesep)
            assert num_kps[0] == 1, "Only support 1 keypoint for now"
            f.write("flip_idx: [0]" + os.linesep)
        f.write("names:" + os.linesep)
        for c in train_cats:
            f.write(f"- {c}" + os.linesep)
        f.write(f"nc: {len(train_cats)}" + os.linesep)
        f.write(f"path: {Path(out_dir).absolute()}" + os.linesep)
        train_rel_path = f"{yolo_train_img_dir.parent.name}/{yolo_train_img_dir.name}"
        f.write(f"train: {train_rel_path}" + os.linesep)
        if yolo_test_img_dir is not None:
            val_rel_path = f"{yolo_test_img_dir.parent.name}/{yolo_test_img_dir.name}"
            f.write(f"val: {val_rel_path}" + os.linesep)

    logger.info(f"Config file saved: {out_config_file}")
    # endregion

    logger.info("Done ✅")


@app.command(help="List all image sizes and counts in a directory recursively")
def list_img_sizes(
    in_dir: str = typer.Argument(..., help="Input directory"),
):
    in_dir_p = Path(in_dir)
    assert in_dir_p.exists(), f"{in_dir_p} does not exist"
    assert in_dir_p.is_dir(), f"{in_dir_p} is not a directory"

    ds = fo.Dataset.from_images_dir(in_dir_p)
    ds.compute_metadata()

    logger.info(f"Found {len(ds)} images in {in_dir_p}")

    # count number of images for each size
    sizes = defaultdict(lambda: 0, {})
    for sample in ds:
        metadata = sample.metadata
        width = metadata.width
        height = metadata.height
        sizes[(width, height)] += 1
    # sort with the most frequent size first
    sizes = dict(sorted(sizes.items(), key=lambda x: x[1], reverse=True))
    for k, v in sizes.items():
        # find one example image for each size
        sample = ds.match({"metadata.width": k[0], "metadata.height": k[1]}).first()
        print(f"Size (w, h) {k}: {v} image(s), e.g., {sample.filepath}")


@app.command(help="Crop images in a directory recursively with ROIs from csv")
def crop_imgs(
    in_dir: str = typer.Argument(..., help="Input directory"),
    roi_csv: str = typer.Argument(..., help="CSV file containing ROIs"),
):
    in_dir_p = Path(in_dir)
    assert in_dir_p.exists(), f"{in_dir_p} does not exist"
    assert in_dir_p.is_dir(), f"{in_dir_p} is not a directory"

    roi_csv_p = Path(roi_csv)
    assert roi_csv_p.exists(), f"{roi_csv_p} does not exist"

    # read ROIs from csv, each image size should have one ROI
    rois = defaultdict(lambda: [], {})
    with open(roi_csv_p, "r") as f:
        for roi in csv.DictReader(f):
            ori_width = int(roi["ori_width"])
            ori_height = int(roi["ori_height"])
            roi_x = int(roi["roi_x"])
            roi_y = int(roi["roi_y"])
            roi_width = int(roi["roi_width"])
            roi_height = int(roi["roi_height"])

            key = (ori_width, ori_height)
            assert key not in rois, f"Duplicate ROI for {key}"
            rois[key] = [roi_x, roi_y, roi_width, roi_height]

    # read and crop images
    # write the cropped images to a new directory
    out_dir_p = in_dir_p.parent / f"{in_dir_p.name}_cropped"
    Path(out_dir_p).mkdir(parents=True, exist_ok=True)
    ds = fo.Dataset.from_images_dir(in_dir_p)
    logger.info(f"Found {len(ds)} images in {in_dir_p}")
    for sample in ds:
        img_path = sample.filepath

        # read and crop the image
        img = cv2.imread(img_path)
        width, height = img.shape[1], img.shape[0]
        roi = rois[(width, height)]
        cropped_img = img[roi[1] : roi[1] + roi[3], roi[0] : roi[0] + roi[2]]

        # keep the original folder structure
        out_img_p = out_dir_p / Path(img_path).relative_to(in_dir_p.absolute())
        # create the subfolder if not exist
        if not out_img_p.parent.exists():
            out_img_p.parent.mkdir(parents=True, exist_ok=True)
        cv2.imwrite(str(out_img_p), cropped_img)
    logger.info(f"Cropped images saved to {out_dir_p}")


@app.command(help="Count num of images without aorta annotations")
def count_n_imgs_no_aorta(
    in_coco_json_p: str = typer.Argument(..., help="Input coco json file"),
    aorta_cat_name: str = typer.Argument("aorta", help="Name of aorta category"),
):
    logger.info(f"Reading {Path(in_coco_json_p).name}...")
    assert Path(in_coco_json_p).exists(), f"{in_coco_json_p} does not exist"
    coco = COCO(in_coco_json_p)

    cats = coco.loadCats(coco.getCatIds())
    cats = sorted(cats, key=lambda x: x["id"], reverse=False)
    # find the category id of aorta
    aorta_cat_id = None
    for cat in cats:
        if cat["name"] == aorta_cat_name:
            aorta_cat_id = cat["id"]
            break
    assert aorta_cat_id is not None, f"Cannot find {aorta_cat_name} in {in_coco_json_p}"
    logger.info(f"Found {aorta_cat_name} with id {aorta_cat_id}")

    n_img_no_aorta = 0
    for img_id in coco.getImgIds():
        anno_ids = coco.getAnnIds(imgIds=img_id)
        annos = coco.loadAnns(anno_ids)
        has_aorta = False
        for anno in annos:
            if anno["category_id"] == aorta_cat_id:
                has_aorta = True
                break
        if not has_aorta:
            n_img_no_aorta += 1
    logger.info(f"Found {n_img_no_aorta} images without {aorta_cat_name}")


@app.command(help="Remove non-aorta annotations from a YOLOv5 dataset")
def keep_only_aorta_labels_in_yolo(
    in_dir: str = typer.Argument(..., help="Input label directory"),
    aorta_class_id: int = typer.Argument(0, help="Class id of aorta"),
):
    txts = list(Path(in_dir).glob("*.txt"))
    logger.info(f"Found {len(txts)} txt files in {in_dir}")
    for txt_p in txts:
        ori_lines, new_lines = [], []
        with open(txt_p, "r") as f:
            ori_lines = f.readlines()
            for line in ori_lines:
                nums = line.split(" ")
                if int(nums[0]) == aorta_class_id:
                    new_lines.append(line)
        with open(txt_p, "w") as new_f:
            new_f.writelines(new_lines)


if __name__ == "__main__":
    app()