Base writer class

This is the base class that other writers inherit from.
It enforces the use of the write() method in all writers.

Bases: ABC

Source code in detection_datasets/writers/base.py

class BaseWriter(ABC):
    def __init__(self, dataset: DetectionDataset, name: str, path: str, move_or_copy_images: str = "copy") -> None:
        """Base class for writing datasets to disk.

        Args:
            dataset: DetectionDataset instance.
            name: Name of the dataset to be created in the "path" directory.
            path: Path to the directory where the dataset will be created.
            move_or_copy_images: Wether to move or copy images from the source
                directory to the directory of the new dataset written to disk.
                Defaults to 'copy'.
        """

        self.dataset = dataset
        self.dataset_dir = os.path.join(path, name)
        self.name = name

        if move_or_copy_images.lower() in ["move", "copy"]:
            self.move_or_copy_images = move_or_copy_images
        else:
            print(f"Incorrect value ({move_or_copy_images}) for move_or_copy_images, defaulting to 'copy'.")
            self.move_or_copy_images = "copy"

    @abstractmethod
    def write(self) -> None:
        """Write the dataset to disk.

        This method is specifc to each format, and need to be implemented in the writer class.
        """

    def do_move_or_copy_image(self, in_file: str, out_file: str) -> None:
        """Move or copy an image.

        Args:
            in_file: Path to the existing image file.
            out_file: Directory where the image will be added.
        """

        if self.move_or_copy_images == "move":
            shutil.move(in_file, out_file)
        else:
            shutil.copyfile(in_file, out_file)

__init__(dataset, name, path, move_or_copy_images='copy')

Base class for writing datasets to disk.

Parameters:

Name	Type	Description	Default
dataset	DetectionDataset	DetectionDataset instance.	required
name	str	Name of the dataset to be created in the "path" directory.	required
path	str	Path to the directory where the dataset will be created.	required
move_or_copy_images	str	Wether to move or copy images from the source directory to the directory of the new dataset written to disk. Defaults to 'copy'.	'copy'

Source code in detection_datasets/writers/base.py

def __init__(self, dataset: DetectionDataset, name: str, path: str, move_or_copy_images: str = "copy") -> None:
    """Base class for writing datasets to disk.

    Args:
        dataset: DetectionDataset instance.
        name: Name of the dataset to be created in the "path" directory.
        path: Path to the directory where the dataset will be created.
        move_or_copy_images: Wether to move or copy images from the source
            directory to the directory of the new dataset written to disk.
            Defaults to 'copy'.
    """

    self.dataset = dataset
    self.dataset_dir = os.path.join(path, name)
    self.name = name

    if move_or_copy_images.lower() in ["move", "copy"]:
        self.move_or_copy_images = move_or_copy_images
    else:
        print(f"Incorrect value ({move_or_copy_images}) for move_or_copy_images, defaulting to 'copy'.")
        self.move_or_copy_images = "copy"

do_move_or_copy_image(in_file, out_file)

Move or copy an image.

Parameters:

Name	Type	Description	Default
in_file	str	Path to the existing image file.	required
out_file	str	Directory where the image will be added.	required

Source code in detection_datasets/writers/base.py

def do_move_or_copy_image(self, in_file: str, out_file: str) -> None:
    """Move or copy an image.

    Args:
        in_file: Path to the existing image file.
        out_file: Directory where the image will be added.
    """

    if self.move_or_copy_images == "move":
        shutil.move(in_file, out_file)
    else:
        shutil.copyfile(in_file, out_file)

write() abstractmethod

Write the dataset to disk.

This method is specifc to each format, and need to be implemented in the writer class.

Source code in detection_datasets/writers/base.py

@abstractmethod
def write(self) -> None:
    """Write the dataset to disk.

    This method is specifc to each format, and need to be implemented in the writer class.
    """

Mmdet writer

This writer saves a dataset to disk in the MMdetection format, that is called "middle format" on the MMdetection documentation.
For more details check Tutorial 2: Customize Datasets on the MMdetection documentation.

Bases: BaseWriter

Source code in detection_datasets/writers/mmdet.py

class MmdetWriter(BaseWriter):

    format = "mmdet"

    def __init__(self, **kwargs) -> None:
        """Initialize the YoloWriter."""

        super().__init__(**kwargs)

    def write(self) -> None:
        """Write the dataset to disk.

        For the MMDET format, the associated steps are:
            1. Create the directories for the images and annotations.
            2. Prepare the data for any given split.
            3. Write the annotation file to disk for each split.
            4. Write the images to disk for each split.
        """

        data = self.dataset.get_data(index="image").reset_index()
        data["bbox"] = [[bbox.to_voc() for bbox in bboxes] for bboxes in data.bbox]

        for split in self.dataset.splits:
            os.makedirs(os.path.join(self.dataset_dir, split, "images"))

            split_data = data[data.split == split]
            dataset = self._make_mmdet_data(split_data)
            self._save_dataset(dataset, split)

    def _make_mmdet_data(self, data_split: pd.DataFrame):

        mmdet_data = []
        source_images = []

        for _, row in data_split.iterrows():
            annotations = {}
            annotations["bboxes"] = row["bbox"]
            annotations["labels"] = row["category_id"]

            data = {}
            data["filename"] = "".join((str(row["image_id"]), ".jpg"))
            data["width"] = row["width"]
            data["height"] = row["height"]
            data["ann"] = annotations

            mmdet_data.append(data)
            source_images.append(row["image_path"])

            dataset = {"mmdet_data": mmdet_data, "source_images": source_images}

        return dataset

    def _save_dataset(self, dataset: Dict[str, List[str]], split: str):
        """Create a new directory and saves the dataset and images."""

        split_path = os.path.join(self.dataset_dir, split)
        mmdet_data = dataset["mmdet_data"]
        source_images = dataset["source_images"]

        # Labels
        file = os.path.join(split_path, "annotation.json")
        with open(file, "w", encoding="utf-8") as f:
            json.dump(mmdet_data, f, ensure_ascii=False, indent=4)

        # Images
        for mmdet_data_image, original_image_path in zip(mmdet_data, source_images):
            out_file = os.path.join(self.dataset_dir, split, "images", mmdet_data_image["filename"])

            self.do_move_or_copy_image(in_file=original_image_path, out_file=out_file)

__init__(**kwargs)

Initialize the YoloWriter.

Source code in detection_datasets/writers/mmdet.py

def __init__(self, **kwargs) -> None:
    """Initialize the YoloWriter."""

    super().__init__(**kwargs)

write()

Write the dataset to disk.

For the MMDET format, the associated steps are: 1. Create the directories for the images and annotations. 2. Prepare the data for any given split. 3. Write the annotation file to disk for each split. 4. Write the images to disk for each split.

Source code in detection_datasets/writers/mmdet.py

def write(self) -> None:
    """Write the dataset to disk.

    For the MMDET format, the associated steps are:
        1. Create the directories for the images and annotations.
        2. Prepare the data for any given split.
        3. Write the annotation file to disk for each split.
        4. Write the images to disk for each split.
    """

    data = self.dataset.get_data(index="image").reset_index()
    data["bbox"] = [[bbox.to_voc() for bbox in bboxes] for bboxes in data.bbox]

    for split in self.dataset.splits:
        os.makedirs(os.path.join(self.dataset_dir, split, "images"))

        split_data = data[data.split == split]
        dataset = self._make_mmdet_data(split_data)
        self._save_dataset(dataset, split)

YOLO writer

This writer saves a dataset to disk in the YOLO format.

Bases: BaseWriter

Write a dataset to a directory in the YOLO format.

Source code in detection_datasets/writers/yolo.py

class YoloWriter(BaseWriter):
    """Write a dataset to a directory in the YOLO format."""

    format = "yolo"

    def __init__(self, **kwargs) -> None:
        """Initialize the YoloWriter."""

        super().__init__(**kwargs)

    def write(self) -> None:
        """Write the dataset to disk.

        For the YOLO format, the associated steps are:
            1. Write the YAML file.
            2. Create the directories for the images and labels.
            3. Write the images and labels.
        """

        data = self.dataset.get_data(index="image").reset_index()
        data["bbox"] = [[bbox.to_yolo() for bbox in bboxes] for bboxes in data.bbox]

        self._write_yaml()

        for split in data.split.unique():
            self._make_dirs(split)

            split_data = data[data.split == split]
            self._write_images_labels(split_data)

    def _write_yaml(self) -> None:
        """Writes the YAML file for the dataset.

        In the YOLO format, this file contains the path to the images, the names of the classes, and the number of
        classes.
        """

        os.makedirs(self.dataset_dir)

        yaml_dataset = {
            "train": f"{self.dataset_dir}/images/train",
            "val": f"{self.dataset_dir}/images/val",
            "test": f"{self.dataset_dir}/images/test",
            "nc": self.dataset.n_categories,
            "names": ", ".join(self.dataset.category_names),
        }

        with open(os.path.join(self.dataset_dir, "dataset.yaml"), "w") as outfile:
            yaml.dump(yaml_dataset, outfile)

    def _make_dirs(self, split: str) -> None:
        """Create the directories (images, labels) for the given split.

        Args:
            split: The split to create the directories for (train, val, test).
        """

        os.makedirs(os.path.join(self.dataset_dir, "images", split))
        os.makedirs(os.path.join(self.dataset_dir, "labels", split))

    def _write_images_labels(self, split_data: pd.DataFrame) -> None:
        """Write the images and labels for a single image.

        Args:
            split_data: The data to write corresponding to a single split.
        """

        for _, row in split_data.iterrows():
            row = row.to_frame().T

            # Images
            in_file = row.image_path.values[0]
            out_file = self._get_filename(row, "images")

            self.do_move_or_copy_image(in_file=in_file, out_file=out_file)

            # Labels
            out_file = self._get_filename(row, "labels")
            data = row.explode(["bbox_id", "category_id", "area", "bbox"])

            with open(out_file, "w") as f:
                for _, r in data.iterrows():
                    labels = " ".join(
                        (str(r.category_id), str(r.bbox[0]), str(r.bbox[1]), str(r.bbox[2]), str(r.bbox[3]))
                    )
                    f.write(labels + "\n")

    def _get_filename(self, row: pd.Series, task: str) -> str:
        """Get the filename for the given row and task.

        Args:
            row: The row of the dataframe to write.
            task: The task to get the filename for (images, labels).

        Returns:
            The filename for the given row and task.

        Raises:
            ValueError: If the task is not images or labels.
        """

        split = row.split.values[0]
        image_id = str(row.image_id.values[0])

        if task == "labels":
            return os.path.join(self.dataset_dir, "labels", split, image_id + ".txt")
        elif task == "images":
            return os.path.join(self.dataset_dir, "images", split, image_id + ".jpg")
        else:
            raise ValueError(f"Task must be either 'lables' or 'images', not {task}")

__init__(**kwargs)

Initialize the YoloWriter.

Source code in detection_datasets/writers/yolo.py

def __init__(self, **kwargs) -> None:
    """Initialize the YoloWriter."""

    super().__init__(**kwargs)

write()

Write the dataset to disk.

For the YOLO format, the associated steps are: 1. Write the YAML file. 2. Create the directories for the images and labels. 3. Write the images and labels.

Source code in detection_datasets/writers/yolo.py

def write(self) -> None:
    """Write the dataset to disk.

    For the YOLO format, the associated steps are:
        1. Write the YAML file.
        2. Create the directories for the images and labels.
        3. Write the images and labels.
    """

    data = self.dataset.get_data(index="image").reset_index()
    data["bbox"] = [[bbox.to_yolo() for bbox in bboxes] for bboxes in data.bbox]

    self._write_yaml()

    for split in data.split.unique():
        self._make_dirs(split)

        split_data = data[data.split == split]
        self._write_images_labels(split_data)

COCO writer

This writer saves a dataset to disk in the COCO format.

Bases: BaseWriter

Source code in detection_datasets/writers/coco.py

class CocoWriter(BaseWriter):

    format = "coco"

    def __init__(self, **kwargs) -> None:
        """Initialize the CocoWriter."""

        super().__init__(**kwargs)

    def write(self) -> None:
        """Write the dataset to disk.

        For the COCO format, the associated steps are:
            1. Write the annotations json files for each split.
            2. Write the images for each split.
        """

        self._write_annotations()
        self._write_images()

    def _write_annotations(self) -> None:
        os.makedirs(os.path.join(self.dataset_dir, "annotations"))

        for split in self.dataset.splits:
            instances = {
                "info": {"description": self.name},
                "licences": [],
                "images": self._get_images(),
                "annotations": self._get_annotations(),
                "categories": self._get_categories(),
            }

            path = os.path.join(self.dataset_dir, "annotations", f"instances_{split}.json")
            with open(path, "w") as file:
                json.dump(instances, file)

    def _get_images(self) -> List[Dict[str, Any]]:
        data = self.dataset.get_data(index="image").reset_index()
        result = []

        for i in range(len(data)):
            result.append(
                {
                    "file_name": str(data.loc[i, "image_path"]).split("/")[-1],
                    "height": data.loc[i, "height"],
                    "width": data.loc[i, "width"],
                    "id": int(data.loc[i, "image_id"]),  # convert from numpy int64
                }
            )

        return result

    def _get_annotations(self) -> List[Dict[str, Any]]:
        data = self.dataset.get_data(index="bbox").reset_index()
        data["bbox"] = [bbox.to_coco() for bbox in data.bbox]
        result = []

        for i in range(len(data)):
            result.append(
                {
                    "area": data.loc[i, "area"],
                    "iscrowd": 0,
                    "image_id": int(data.loc[i, "image_id"]),
                    "bbox": data.loc[i, "bbox"],
                    "category_id": data.loc[i, "category_id"],
                    "id": int(data.loc[i, "bbox_id"]),
                }
            )

        return result

    def _get_categories(self) -> List[Dict[str, Any]]:
        data = self.dataset.categories.copy().reset_index()
        result = []

        for i in range(len(data)):
            result.append(
                {
                    "id": int(data.loc[i, "category_id"]),
                    "name": data.loc[i, "category"],
                }
            )

        return result

    def _write_images(self) -> None:
        data = self.dataset.get_data(index="image").reset_index()

        for split in self.dataset.splits:
            os.makedirs(os.path.join(self.dataset_dir, split))
            split_data = data[data.split == split]

            for _, row in split_data.iterrows():
                row = row.to_frame().T

                in_file = row.image_path.values[0]
                out_file = os.path.join(self.dataset_dir, split, str(row["image_id"].values[0]) + ".jpg")

                self.do_move_or_copy_image(in_file=in_file, out_file=out_file)

__init__(**kwargs)

Initialize the CocoWriter.

Source code in detection_datasets/writers/coco.py

def __init__(self, **kwargs) -> None:
    """Initialize the CocoWriter."""

    super().__init__(**kwargs)

write()

Write the dataset to disk.

For the COCO format, the associated steps are: 1. Write the annotations json files for each split. 2. Write the images for each split.

Source code in detection_datasets/writers/coco.py

def write(self) -> None:
    """Write the dataset to disk.

    For the COCO format, the associated steps are:
        1. Write the annotations json files for each split.
        2. Write the images for each split.
    """

    self._write_annotations()
    self._write_images()