|
|
"""
|
|
|
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
|
|
|
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
|
|
|
"""
|
|
|
|
|
|
import torch
|
|
|
import torch.utils.data as data
|
|
|
|
|
|
|
|
|
class DetDataset(data.Dataset):
|
|
|
def __getitem__(self, index):
|
|
|
img, target = self.load_item(index)
|
|
|
if self.transforms is not None:
|
|
|
img, target, _ = self.transforms(img, target, self)
|
|
|
return img, target
|
|
|
|
|
|
def load_item(self, index):
|
|
|
raise NotImplementedError(
|
|
|
"Please implement this function to return item before `transforms`."
|
|
|
)
|
|
|
|
|
|
def set_epoch(self, epoch) -> None:
|
|
|
self._epoch = epoch
|
|
|
|
|
|
@property
|
|
|
def epoch(self):
|
|
|
return self._epoch if hasattr(self, "_epoch") else -1
|
|
|
|
