updated to use data_name_list

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)