diff --git a/model/src/dataloader.py b/model/src/dataloader.py index d5cc0e5..77c433e 100644 --- a/model/src/dataloader.py +++ b/model/src/dataloader.py @@ -15,7 +15,6 @@ from torchvision.transforms.functional import resize from torchvision.transforms.functional import rotate from core.database import connect -from core.database import list_names # resnet18 original normalization values RESNET_NORMALIZE_MEAN = [0.485, 0.456, 0.406] @@ -25,12 +24,14 @@ RESNET_NORMALIZE_STD = [0.229, 0.224, 0.225] class VCDADataset(Dataset): def __init__( self, + data_name_list: list[str], do_augment: bool = False, random_annotations: bool = False, normalize_mean: list[float] = RESNET_NORMALIZE_MEAN, normalize_std: list[float] = RESNET_NORMALIZE_STD, ): super().__init__() + self.data_name_list = data_name_list self.do_augment = do_augment self.random_annotations = random_annotations self.normalize_mean = normalize_mean @@ -48,12 +49,6 @@ class VCDADataset(Dataset): # connect to database collection, _, _ = connect() self.collection = collection - # load data names - has_annotation = False if self.random_annotations else True - self.data_name_list = list_names( - collection=self.collection, - has_annotation=has_annotation, - ) def __len__(self): return len(self.data_name_list)