initial ppush

Dockerfile

@@ -0,0 +1,49 @@
+FROM nvcr.io/nvidia/pytorch:23.03-py3
+# FROM nvcr.io/nvidia/pytorch:24.02-py3
+
+ARG HOME_PATH="/home"
+WORKDIR ${HOME_PATH}
+
+RUN pip3 install --upgrade pip wheel
+RUN pip3 install azure-storage-blob azure-identity
+
+# supervision
+RUN git clone https://github.com/roboflow/supervision.git && \
+    cd supervision && \
+    grep -v "^opencv-python-headless" pyproject.toml > tmp.toml && \
+    mv tmp.toml pyproject.toml && \
+    pip3 install --no-cache -e .
+
+# ultralytics
+ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Arial.Unicode.ttf /root/.config/Ultralytics/
+RUN git clone https://github.com/ultralytics/ultralytics && \
+    cd ultralytics && \
+    grep -v "opencv-python\|openvino-dev" pyproject.toml > tmp.toml && mv tmp.toml pyproject.toml && \
+    pip3 install "opencv-python-headless<4.7" "opencv-contrib-python<4.7" "opencv-contrib-python-headless<4.7" "albumentations<1.4.0" && \
+    pip3 install .
+
+# download dataset
+ARG CVAT_URL
+ARG CVAT_ORG
+ARG CVAT_TASKS_YAML
+ARG TRAIN_HP_YAML
+ARG PYPREPROCESS
+
+COPY . .
+# COPY AIEM/trainer /home/trainer
+# COPY ${CVAT_TASKS_YAML} ${CVAT_TASKS_YAML}
+# COPY ${TRAIN_HP_YAML} ${TRAIN_HP_YAML}
+
+ENV APP_PYPREPROCESS=${PYPREPROCESS}
+ENV APP_CVAT_TASKS_YAML=${CVAT_TASKS_YAML}
+ENV APP_HOME=${HOME_PATH}
+ENV APP_TRAIN_HP_YAML=${TRAIN_HP_YAML}
+
+RUN cd AIEM/trainer && \
+    python3 utils/download_cvatdata.py \
+        "$CVAT_URL" \
+        "$CVAT_ORG"
+RUN cd /data && \
+    rm -rf *.zip
+
+ENTRYPOINT ["python3", "AIEM/trainer/train_yolov8.py"]

Dockerfile.x86.yolov8_trainer

@@ -0,0 +1,49 @@
+FROM nvcr.io/nvidia/pytorch:23.03-py3
+# FROM nvcr.io/nvidia/pytorch:24.02-py3
+
+ARG HOME_PATH="/home"
+WORKDIR ${HOME_PATH}
+
+RUN pip3 install --upgrade pip wheel
+RUN pip3 install azure-storage-blob azure-identity
+
+# supervision
+RUN git clone https://github.com/roboflow/supervision.git && \
+    cd supervision && \
+    grep -v "^opencv-python-headless" pyproject.toml > tmp.toml && \
+    mv tmp.toml pyproject.toml && \
+    pip3 install --no-cache -e .
+
+# ultralytics
+ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Arial.Unicode.ttf /root/.config/Ultralytics/
+RUN git clone https://github.com/ultralytics/ultralytics && \
+    cd ultralytics && \
+    grep -v "opencv-python\|openvino-dev" pyproject.toml > tmp.toml && mv tmp.toml pyproject.toml && \
+    pip3 install "opencv-python-headless<4.7" "opencv-contrib-python<4.7" "opencv-contrib-python-headless<4.7" "albumentations<1.4.0" && \
+    pip3 install .
+
+# download dataset
+ARG CVAT_URL
+ARG CVAT_ORG
+ARG CVAT_TASKS_YAML
+ARG TRAIN_HP_YAML
+ARG PYPREPROCESS
+
+COPY . .
+# COPY AIEM/trainer /home/trainer
+# COPY ${CVAT_TASKS_YAML} ${CVAT_TASKS_YAML}
+# COPY ${TRAIN_HP_YAML} ${TRAIN_HP_YAML}
+
+ENV APP_PYPREPROCESS=${PYPREPROCESS}
+ENV APP_CVAT_TASKS_YAML=${CVAT_TASKS_YAML}
+ENV APP_HOME=${HOME_PATH}
+ENV APP_TRAIN_HP_YAML=${TRAIN_HP_YAML}
+
+RUN cd AIEM/trainer && \
+    python3 utils/download_cvatdata.py \
+        "$CVAT_URL" \
+        "$CVAT_ORG"
+RUN cd /data && \
+    rm -rf *.zip
+
+ENTRYPOINT ["python3", "AIEM/trainer/train_yolov8.py"]

README.md

@@ -1,10 +1,34 @@
----
-title: AIEM
-emoji: 🐠
-colorFrom: blue
-colorTo: blue
-sdk: docker
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# AIEM
+
+AI Edge Management
+
+**TODO**: introduce segmentation env variable 
+
+AIEM repo can be seen as the core shared across all the projects that require an AI model to be trained or to run an inference server. It talks to the rest of the project-specific repos by means of, e.g., a GitHub Actions workflow. It contains Dockerfiles for different architectures and for different purposes. For example: training a YoloV8 model in an x86 architecture (*Dockerfile.x86.yolov8_trainer*). 
+
+## Structure
+
+The structure of the project:
+
+```bash
+.
+├── docker
+│   ├── Dockerfile.x86.yolov8_trainer
+│   └── scripts
+│       └── docker_build.sh
+├── README.md
+├── runner
+│   └── README.md
+└── trainer
+    ├── README.md
+    ├── train_yolov8.py
+    └── utils
+        ├── cvat_dataset.py
+        ├── download_cvatdata.py
+        ├── merge_cocos.py
+        ├── path_utils.py
+        ├── unzip_datasets.py
+        └── yolo_labels.py
+```
+
+- **Download data** (*trainer/utils/download_cvatdata.py*). Main script to download the dataset into the docker container. It reads from project-specific YAML file with the tasks to download from CVAT, preprocess the data and get the workspace ready for the model be able to be trained.

docker/scripts/docker_build.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+
+CONTAINER=$1
+DOCKERFILE=$2
+
+shift 
+shift
+
+echo "Building $CONTAINER container..."
+
+docker build --network=host -t $CONTAINER -f $DOCKERFILE "$@" .

runner/README.md

@@ -0,0 +1 @@
+To be designed. It has to do with project-wise running models. 

trainer/README.md

@@ -0,0 +1 @@
+This folder reads from a config file specific to a project. This repo must be sym-linked inside the project folder.

trainer/train_yolov8.py

@@ -0,0 +1,146 @@
+import argparse
+import os
+import yaml
+import shutil 
+import datetime
+import numpy as np
+import pandas as pd 
+import yaml
+from azure.storage.blob import BlobServiceClient
+from pathlib import Path
+from sklearn.model_selection import KFold
+from collections import Counter
+from ultralytics import YOLO
+from utils.path_utils import *
+
+STORAGE_ACCOUNT_KEY  = "mhqTCNmdIgsnvyFnfv0r2JKfs8iG//5YVnphCq336XNxhyI72brMy6lP88I9XKVya/G9ZlAAMoNd+AStsXFe0Q=="
+STORAGE_ACCOUNT_NAME = "camtagstoreaiem"
+CONNECTION_STRING    = "DefaultEndpointsProtocol=https;AccountName=camtagstoreaiem;AccountKey=mhqTCNmdIgsnvyFnfv0r2JKfs8iG//5YVnphCq336XNxhyI72brMy6lP88I9XKVya/G9ZlAAMoNd+AStsXFe0Q==;EndpointSuffix=core.windows.net"
+CONTAINER_NAME       = "upload"
+
+# Get YAML file containing the training hyperparameters
+HOME = os.getenv("APP_HOME")
+APP_TRAIN_HP_YAML = os.path.join(HOME, os.getenv("APP_TRAIN_HP_YAML"))
+
+def azure_upload(local_fname, blob_fname, overwrite=True):
+    blob_service_client = BlobServiceClient.from_connection_string(CONNECTION_STRING)
+    blob_client = blob_service_client.get_blob_client(
+        container = CONTAINER_NAME,
+        blob = blob_fname
+    )
+    with open(local_fname, "rb") as data:
+        blob_client.upload_blob(data, overwrite=overwrite)
+
+
+if __name__ == "__main__":
+    with open(APP_TRAIN_HP_YAML, "r") as f:
+        y = yaml.safe_load(f)
+        KSPLIT     = y['ksplit']
+        EPOCHS     = y['epochs']
+        MODEL      = y['model']
+        DATA_PATH  = y['data_path']
+        BATCH_SIZE = y['batch_size']
+
+    # coco
+    coco_dataset_path = Path(DATA_PATH)
+    coco_dict = read_coco_json(coco_dataset_path / "merged.json")
+    
+    classes = {cat['id']-1: cat['name'] for cat in coco_dict['categories']}
+    cls_idx = sorted(classes.keys())
+
+    labels = sorted((coco_dataset_path / "labels").rglob("*.txt"))
+    indx = [l.stem for l in labels]
+    labels_df = pd.DataFrame([], columns=cls_idx, index=indx)
+    
+    for label in labels:
+        label_counter = Counter()
+        with open(label, 'r') as lf:
+            lines = lf.readlines()
+        
+        for l in lines:
+            label_counter[int(l.split(' ')[0])] += 1
+        labels_df.loc[label.stem] = label_counter
+
+    labels_df = labels_df.fillna(0.0)
+
+    # KFOLD
+    kf = KFold(
+        n_splits = KSPLIT,
+        shuffle = True,
+        random_state = 42
+    )
+    kfolds = list(kf.split(labels_df))
+
+    folds = [f'split_{n}' for n in range(1, KSPLIT + 1)]
+    folds_df = pd.DataFrame(index=indx, columns=folds)
+    for idx, (train, val) in enumerate(kfolds, start=1):
+        folds_df[f'split_{idx}'].loc[labels_df.iloc[train].index] = 'train'
+        folds_df[f'split_{idx}'].loc[labels_df.iloc[val].index] = 'val'
+
+    # check distributions. balanced?
+    fold_lbl_distrb = pd.DataFrame(index=folds, columns=cls_idx)
+    for n, (train_indices, val_indices) in enumerate(kfolds, start=1):
+        train_totals = labels_df.iloc[train_indices].sum()
+        val_totals = labels_df.iloc[val_indices].sum()
+
+        ratio = val_totals / (train_totals + 1E-7)
+        fold_lbl_distrb.loc[f'split_{n}'] = ratio
+
+    # datasets for each fold
+    save_path = Path(coco_dataset_path / f'{datetime.date.today().isoformat()}_{KSPLIT}-Fold_Cross-val')
+    save_path.mkdir(parents=True, exist_ok=True)
+
+    suffix = sorted((coco_dataset_path / 'images').rglob("*.*"))[0].suffix
+    images = [coco_dataset_path / "images" / l.with_suffix(suffix).name for l in labels]
+    ds_yamls = []
+    
+    for split in folds_df.columns:
+        # create directories
+        split_dir = save_path / split
+        split_dir.mkdir(parents=True, exist_ok=True)
+        (split_dir / 'train' / 'images').mkdir(parents=True, exist_ok=True)
+        (split_dir / 'train' / 'labels').mkdir(parents=True, exist_ok=True)
+        (split_dir / 'val' / 'images').mkdir(parents=True, exist_ok=True)
+        (split_dir / 'val' / 'labels').mkdir(parents=True, exist_ok=True)
+
+        # create yaml files
+        dataset_yaml = split_dir / f'{split}_dataset.yaml'
+        ds_yamls.append(dataset_yaml)
+
+        with open(dataset_yaml, 'w') as ds_y:
+            yaml.safe_dump({
+                'path' : split_dir.resolve().as_posix(),
+                'train': 'train',
+                'val'  : 'val',
+                'names': classes
+            }, ds_y)
+
+    for image, label in zip(images, labels):
+        for split, k_split in folds_df.loc[image.stem].items():
+            # destination directory 
+            img_to_path = save_path / split / k_split / 'images'
+            lbl_to_path = save_path / split / k_split / 'labels'
+
+            # copy image and label file to new directory
+            shutil.copy(image, img_to_path / image.name)
+            shutil.copy(label, lbl_to_path / label.name)
+
+    folds_df.to_csv(save_path / "kfold_datasplit.csv")
+    fold_lbl_distrb.to_csv(save_path / "kfold_label_distributions.csv")
+
+    model = YOLO(MODEL)
+    
+    for k in range(KSPLIT):
+        dataset_yaml = ds_yamls[k]
+        model.train(
+            data = dataset_yaml, 
+            epochs = EPOCHS,
+            batch = BATCH_SIZE,
+            plots = False
+        )
+
+    # azure upload
+    flag = '2' * (KSPLIT - 1)
+    local_fname = f'runs/detect/train{flag}/weights/best.pt'
+    blob_fname = f"kohberg/host_train_{MODEL}"
+    azure_upload(local_fname, blob_fname, overwrite=True)

trainer/utils/cvat_dataset.py

@@ -0,0 +1,108 @@
+import os
+import sys
+import requests
+import shutil
+import time 
+from pathlib import Path
+from tqdm.auto import tqdm
+
+class CVATDataset:
+    def __init__(self, cvat_url, org, task_ids, headers=None, params=None, names=None, dest_folder=None):
+        """
+        Connects to serverless CVAT to download datasets. 
+
+        Args:
+            cvat_url    (str) : CVAT base URL where the server is loaded.
+            org         (str) : organization we are working with, e.g.: 'bulow'
+            task_ids    (list): list with the task IDs inside CVAT. 
+            params      (dict): query parameters. 
+            names       (dict): dict where the keys are the task id and values 
+                                the names of the local files.
+            dest_folder (str) : destination folder of the zip files.
+        
+        Returns:
+            Content ZIP file containing JSON coco annotations and the images. 
+        """
+        self.cvat_url = cvat_url
+        self.org = org
+        self.task_ids = task_ids
+        self.dest_folder = dest_folder
+        self.names_dict = names
+        if self.names_dict is not None:
+            assert all([id_ in self.names_dict.keys() for id_ in self.task_ids]), \
+                "The keys in names do not match the task IDs."
+
+        self.headers = headers
+        if self.headers is None:
+            # FIXME: avoid hardcoded authorization. 
+            self.headers = {"Authorization": "Basic ZGphbmdvOlMwbHNraW4xMjM0IQ=="}
+        
+        self.params = params
+        if self.params is None:
+            self.params = {
+                "format"  : "COCO 1.0",
+                "action"  : "download",
+                "location": "local",
+                "org"     : self.org
+            }
+
+    @staticmethod
+    def countdown_clock(waiting_time):
+        t0 = time.monotonic()
+        while time.monotonic() - t0 < waiting_time:
+            remaining_time = waiting_time - (time.monotonic() - t0)
+            mins, secs = divmod(int(remaining_time), 60)
+            sys.stdout.write("\r")
+            sys.stdout.write(f"{mins:02d}:{secs:02d}")
+            sys.stdout.flush()
+            time.sleep(1)
+        sys.stdout.write("\n")
+    
+    def _get_dataset(self, endpoint):
+        response = requests.get(
+            endpoint,
+            headers = self.headers,
+            params = self.params,
+            stream = True
+        )
+        return response
+    
+    def _download_task(self, task_id: int, fname: str):
+        """ Downloads dataset linked to a task. """
+        endpoint = f"{self.cvat_url}/api/tasks/{task_id}/dataset"
+        r = self._get_dataset(endpoint)
+        while r.status_code != 200:
+            if r.status_code == 202:
+                print(f"  Status code {r.status_code}: server processing request")
+                self.countdown_clock(10)
+            else:
+                print(f"  Status code {r.status_code}: connection error")
+                self.countdown_clock(30)
+            r = self._get_dataset(endpoint)
+        
+        print(f"  Status code {r.status_code}: request is ready")
+        total_length = int(r.headers.get("Content-Length"))
+        with tqdm.wrapattr(r.raw, "read", total=total_length, desc="") as raw:
+            with open(fname, "wb") as file:
+                shutil.copyfileobj(raw, file)
+
+    def download_tasks(self):
+        """ Download all the tasks passed as input. """
+        for task_id in self.task_ids:
+            name_label = task_id
+            if self.names_dict is not None:
+                name_label = self.names_dict[task_id]
+            fname = f"dataset_{name_label}.zip"
+            if self.dest_folder is not None:
+                self.dest_folder = Path(self.dest_folder)
+                self.dest_folder.mkdir(exist_ok=True, parents=True)
+            fname = (self.dest_folder / fname).resolve().as_posix()
+
+            if os.path.exists(fname):
+                print(f"File {fname} already exists.")
+                continue
+            
+            print(f"\nDownloading task {task_id}, with fname {fname}")
+            self._download_task(task_id, fname)
+
+    # TODO: implement unzip function for the tasks

trainer/utils/download_cvatdata.py

@@ -0,0 +1,98 @@
+"""
+This script reads from a YAML file and downloads data from CVAT.
+"""
+
+import os
+import argparse
+import subprocess 
+import shutil
+import yaml
+from pathlib import Path
+from cvat_dataset import CVATDataset
+from merge_cocos import merge
+from yolo_labels import get_yolo_labels
+
+HOME = os.getenv("APP_HOME")
+CVAT_TASKS = os.path.join(HOME, os.getenv("APP_CVAT_TASKS_YAML"))
+PYPREPROCESS = os.getenv("APP_PYPREPROCESS")
+import sys
+sys.path.append(HOME)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        'cvat_url', 
+        type = str, 
+        help = 'cvat url'
+    )
+    parser.add_argument(
+        'cvat_org',
+        type = str, 
+        help = 'cvat organization'
+    )
+    parser.add_argument(
+        '-odir', '--output_dir', 
+        type = str, 
+        help = "path to download directory",
+        default = "/data"
+    )
+    args = parser.parse_args()
+
+    with open(CVAT_TASKS, "r") as f:
+        y = yaml.safe_load(f)
+        TASK_IDS = y["task_ids"]
+        NAMES = None
+        if "names" in y:
+            NAMES = y["names"]
+
+    data_folder = Path(args.output_dir)
+    data_folder.mkdir(parents=True, exist_ok=True)
+
+    CVAT = CVATDataset(
+        args.cvat_url,
+        args.cvat_org,
+        TASK_IDS,
+        names = NAMES, 
+        dest_folder = data_folder
+    )
+    CVAT.download_tasks()
+
+    paths2imgs = []
+    paths2json = []
+    paths2dirs = []
+    for dataset in data_folder.rglob("*.zip"):
+        dir_name = dataset.parent / dataset.stem
+        paths2dirs.append(dir_name)
+        paths2imgs.append(dir_name / "images")
+        paths2json.append(dir_name / "annotations" / "instances_default.json")
+        if dir_name.exists():
+            continue
+        subprocess.call(['unzip', '-o', dataset, '-d', dir_name])
+    
+    if PYPREPROCESS == 'true':
+        # looks for the py script called: trainer_files/preprocess.py
+        # this script is characteristic to the project
+        from trainer_files.preprocess import preprocess_cvat
+        paths2json, paths2imgs = preprocess_cvat(paths2dirs)
+
+    # TODO: add debugging / assert script to make sure preprocess is done correctly
+    
+    # merge everything into a single json file
+    if len(paths2json) > 1:
+        merge(
+            paths2json, paths2imgs, data_folder / 'merged_cocos', 'merged', verbose=True
+        )
+    else:
+        json_file = Path(paths2json[0])
+        shutil.copy(
+            json_file.as_posix(), 
+            (json_file.parents[1] / 'merged.json').as_posix()
+        )
+        shutil.move(
+            json_file.parents[1].as_posix(), 
+            (data_folder / 'merged_cocos').as_posix()
+        )
+
+    # yolo format - labels
+    path2json = data_folder / 'merged_cocos' / 'merged.json'
+    get_yolo_labels(path2json, use_segment=False)

trainer/utils/merge_cocos.py

@@ -0,0 +1,98 @@
+import os
+import glob
+from pathlib import Path
+from datetime import date
+from collections import defaultdict 
+from warnings import warn 
+
+from path_utils import *
+
+def merge_cats_get_id(cats, this_cat):
+    cat_nms = [c['name'] for c in cats]
+    if this_cat['name'] not in cat_nms:
+        this_cat['id'] = len(cats) + 1
+        cats.append(this_cat)
+        return this_cat["id"]
+    else:
+        return this_cat["id"]
+    
+
+def filter_images(images, annotations):
+    img_ids_from_anns = [ann['image_id'] for ann in annotations]
+    images_ = [
+        img_info for img_info in images if img_info['id'] in img_ids_from_anns
+    ]
+    return images_
+
+
+def merge(jsons, img_roots, output_dir, output_nm="merged", verbose=True):
+    assert len(jsons) == len(img_roots)
+
+    out_dir_path = Path(output_dir)
+    out_imgs_dir_path = out_dir_path / "images"
+    
+    merged_img_id_state = 1
+    merged_ann_id_state = 1
+    merged_names = []
+    merged_dict = {
+        "info"       : {"description": "", "data_created": f"{date.today():%Y/%m/%d}"},
+        "annotations": [],
+        "categories" : [],
+        "images"     : [] 
+    }
+    for i, (json_path, imgs_dir_path) in enumerate(zip(jsons, img_roots)):
+        coco_dict = read_coco_json(json_path)
+        dataset_name = get_setname(json_path)
+        merged_names.append(dataset_name)
+
+        # categories
+        cat_id_old2new = {}
+        for cat in coco_dict['categories']:
+            old_cat_id = cat['id']
+            new_cat_id = merge_cats_get_id(merged_dict['categories'], cat)
+            cat_id_old2new[old_cat_id] = new_cat_id
+        
+        # images
+        coco_dict['images'] = filter_images(
+            coco_dict['images'], coco_dict['annotations']
+        )
+        img_id_old2new = {}
+        for img in coco_dict['images']:
+            img_id_old2new[img["id"]] = merged_img_id_state
+            img["id"] = merged_img_id_state
+
+            old_img_path = Path(imgs_dir_path) / img['file_name']
+            img['file_name'] = dataset_name + "_" + img['file_name']
+            new_img_path = out_imgs_dir_path / img['file_name']
+            assure_copy(old_img_path, new_img_path)
+
+            merged_img_id_state += 1
+            merged_dict['images'].append(img)
+
+        # annotations
+        for ann in coco_dict['annotations']:
+            ann['id'] = merged_ann_id_state
+            ann['image_id'] = img_id_old2new[ann['image_id']]
+            ann['category_id'] = cat_id_old2new[ann['category_id']]
+
+            merged_ann_id_state += 1
+            merged_dict['annotations'].append(ann)
+
+    merged_dict["info"]["description"] = "+".join(merged_names)
+
+    out_json = out_dir_path / f"{output_nm}.json"
+    write_json(out_json, merged_dict)
+
+    if verbose:
+        print(f"Number of images: {len(merged_dict['images'])}")
+        print(f"Number of annotations: {len(merged_dict['annotations'])}")
+
+
+if __name__ == '__main__':
+    paths2images = []
+    paths2json = []
+    for dataset in glob.glob("dataset_*"):
+        paths2images.append(os.path.join(dataset, "images"))
+        paths2json.append(os.path.join(dataset, "annotations/instances_default.json"))
+
+    merge(paths2json, paths2images, './merged_cocos', 'merged', verbose=True)

trainer/utils/path_utils.py

@@ -0,0 +1,48 @@
+import json 
+import filecmp
+from pathlib import Path
+from shutil import copy
+
+
+def read_json(json_path):
+    with open(json_path, "r") as f:
+        d = json.load(f)
+    return d
+
+
+def write_json(json_path, dic):
+    with open(json_path, "w") as f:
+        json.dump(dic, f)
+    print(f"Wrote json to {json_path}")
+
+
+def get_setname(json_path):
+    json_path_ = Path(json_path)
+    dataset_nm = json_path_.parent.parts[-2]
+    print(f"Processing {dataset_nm} (name derived from json path)")
+    return dataset_nm
+
+
+def read_coco_json(coco_json):
+    coco_dict = read_json(coco_json)
+    return coco_dict
+
+
+def assure_copy(src, dst):
+    assert Path(src).is_file()
+    if Path(dst).is_file() and filecmp.cmp(src, dst, shallow=True):
+        return 
+    Path(dst).parent.mkdir(exist_ok=True, parents=True)
+    copy(src, dst)
+
+
+def path(str_path, is_dir=False, mkdir=False):
+    path_ = Path(str_path)
+    if is_dir:
+        if mkdir:
+            path_.mkdir(parents=True, exist_ok=True)
+        assert path_.is_dir(), path_
+    else:
+        assert path_.is_file(), path_
+    return path_
+

trainer/utils/unzip_datasets.py

@@ -0,0 +1,7 @@
+import subprocess 
+import glob
+
+if __name__ == "__main__":
+    for dataset in glob.glob("*.zip"):
+        dir_name = dataset.split(".")[0]
+        subprocess.call(['unzip', '-o', dataset, '-d', dir_name])

trainer/utils/yolo_labels.py

@@ -0,0 +1,128 @@
+import numpy as np
+import json
+from pathlib import Path, PosixPath
+from pycocotools.coco import COCO
+
+def min_index(arr1, arr2):
+    """
+    Find a pair of indexes with the shortest distance.
+
+    Args:
+        arr1: (N, 2).
+        arr2: (M, 2).
+    Return:
+        a pair of indexes (tuple)
+    """
+    dis = ((arr1[:, None, :] - arr2[None, :, :]) ** 2).sum(-1)
+    return np.unravel_index(np.argmin(dis, axis=None), dis.shape)
+
+
+def merge_multi_segment(segments):
+    """
+    Merge multi segments to one list.
+    Find coordinates with min distance between each segment,
+    then connect these coordinates with one thin line to merge all
+    segments into one.
+
+    Args:
+        segments (List(List)): original segmentations in coco's json file
+            like [segmentation1, segmentation2, ...], where
+            each segmentation is a list of coordinates
+    """
+    s = []
+    segments = [np.array(i).reshape(-1,2) for i in segments]
+    idx_list = [[] for _ in range(len(segments))]
+
+    # record the indexes with the min distance between each segment
+    for i in range(1, len(segments)):
+        idx1, idx2 = min_index(segments[i - 1, segments[i]])
+        idx_list[i - 1].append(idx1)
+        idx_list[i].append(idx2)
+
+    # use two round to connect all the segments
+    for k in range(2):
+        # forward connection
+        if k == 0:
+            for i, idx in enumerate(idx_list):
+                # middle segments have two indexes
+                # reverse the index of middle segments
+                if len(idx) == 2 and idx[0] > idx[1]:
+                    idx = idx[::-1]
+                    segments[i] = segments[i][::-1, :]
+
+                segments[i] = np.roll(segments[i], -idx[0], axis=0)
+                segments[i] = np.concatenate([segments[i], segments[i][:1]])
+
+                # deal with the first segment and the last one
+                if i in [0, len(idx_list) - 1]:
+                    s.append(segments[i])
+                else:
+                    idx = [0, idx[1] - idx[0]]
+                    s.append(segments[i][idx[0]:idx[1] + 1])
+
+        else:
+            for i in range(len(idx_list) - 1, -1, -1):
+                if i not in [0, len(idx_list) - 1]:
+                    idx = idx_list[i]
+                    nidx = abs(idx[1] - idx[0])
+                    s.append(segments[i][nidx:])
+    return s
+
+
+def get_yolo_labels(path2json, use_segment=False):
+    if not isinstance(path2json, PosixPath):
+        path2json = Path(path2json)
+    path2labels = path2json.parents[0] / "labels"
+    path2labels.mkdir(parents=True, exist_ok=True)
+
+    coco = COCO(path2json)
+
+    img2anns = {}
+    for ann in coco.dataset['annotations']:
+        img_id = ann['image_id']
+        if img_id not in img2anns:
+            img2anns[img_id] = [ann]
+        else:
+            img2anns[img_id].append(ann)
+
+    id2img = {img["id"]: img for img in coco.dataset["images"]}
+
+    for img_id, anns in img2anns.items():
+        img = id2img[img_id]
+        h, w, f = img['height'], img['width'], img['file_name']
+
+        bboxes = []
+        segments = []
+        for ann in anns:
+            if ann['iscrowd']:
+                continue
+            # coco box format: [top left x, top left y, width, height]
+            box = np.array(ann['bbox'], dtype=np.float64)
+            box[:2] += box[2:] / 2  # center coordinates
+            box[[0, 2]] /= w  # normalize x
+            box[[1, 3]] /= h  # normalize y
+            if box[2] <= 0 or box[3] <= 0:
+                continue
+
+            cls = ann['category_id'] - 1
+            box = [cls] + box.tolist()
+            if box not in bboxes:
+                bboxes.append(box)
+
+            # segmentation?
+            if use_segment:
+                if len(ann['segmentation']) > 1:
+                    s = merge_multi_segment(ann['segmentation'])
+                    s = (np.concatenate(s, axis=0) / np.array([w, h])).reshape(-1).tolist()
+                else:
+                    s = [j for i in ann['segmentation'] for j in i]  # all segments concatenated
+                    s = (np.array(s).reshape(-1, 2) / np.array([w, h])).reshape(-1).tolist()
+                s = [cls] + s
+                if s not in segments:
+                    segments.append(s)
+
+        # write
+        with open((path2labels / f).with_suffix('.txt'), 'a') as file:
+            for i in range(len(bboxes)):
+                line = *(segments[i] if use_segment else bboxes[i]),
+                file.write(('%g ' * len(line)).rstrip() % line + '\n')