Spaces:
Runtime error
Runtime error
File size: 1,113 Bytes
153628e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
# Copyright (C) 2021-2024, Mindee.
# This program is licensed under the Apache License 2.0.
# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
import os
from typing import Any, List, Tuple
import numpy as np
from .datasets import AbstractDataset
__all__ = ["OrientationDataset"]
class OrientationDataset(AbstractDataset):
"""Implements a basic image dataset where targets are filled with zeros.
>>> from doctr.datasets import OrientationDataset
>>> train_set = OrientationDataset(img_folder="/path/to/images")
>>> img, target = train_set[0]
Args:
----
img_folder: folder with all the images of the dataset
**kwargs: keyword arguments from `AbstractDataset`.
"""
def __init__(
self,
img_folder: str,
**kwargs: Any,
) -> None:
super().__init__(
img_folder,
**kwargs,
)
# initialize dataset with 0 degree rotation targets
self.data: List[Tuple[str, np.ndarray]] = [(img_name, np.array([0])) for img_name in os.listdir(self.root)]
|