File size: 1,217 Bytes
afc2161 |
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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
from abc import ABC, abstractmethod
from typing import Any
from torch.utils.data import Dataset
from ellipse_rcnn.utils.types import (
TargetDict,
CollatedBatchType,
UncollatedBatchType,
)
def collate_fn(batch: UncollatedBatchType) -> CollatedBatchType:
"""
Collate function for the :class:`DataLoader`.
Parameters
----------
batch:
A batch of data.
"""
return tuple(zip(*batch)) # type: ignore
class EllipseDatasetBase(ABC, Dataset):
@abstractmethod
def load_image(self, index: int) -> Any:
"""
Load the image for the given index.
Parameters
----------
index:
The index of the image.
Returns
-------
image:
The raw image.
"""
pass
@abstractmethod
def load_target_dict(self, index: int) -> TargetDict:
"""
Load the target dict for the given index.
Parameters
----------
index:
The index of the target dict.
Returns
-------
target_dict:
The target dictionary.
"""
pass
@abstractmethod
def __len__(self) -> int:
pass
|