updated to use data_name_list

This commit is contained in:
Brian Bjarke Jensen
2024-03-12 18:45:35 +01:00
parent 1065a35cf0
commit ec48879327
+2 -7
View File
@@ -15,7 +15,6 @@ from torchvision.transforms.functional import resize
from torchvision.transforms.functional import rotate from torchvision.transforms.functional import rotate
from core.database import connect from core.database import connect
from core.database import list_names
# resnet18 original normalization values # resnet18 original normalization values
RESNET_NORMALIZE_MEAN = [0.485, 0.456, 0.406] 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): class VCDADataset(Dataset):
def __init__( def __init__(
self, self,
data_name_list: list[str],
do_augment: bool = False, do_augment: bool = False,
random_annotations: bool = False, random_annotations: bool = False,
normalize_mean: list[float] = RESNET_NORMALIZE_MEAN, normalize_mean: list[float] = RESNET_NORMALIZE_MEAN,
normalize_std: list[float] = RESNET_NORMALIZE_STD, normalize_std: list[float] = RESNET_NORMALIZE_STD,
): ):
super().__init__() super().__init__()
self.data_name_list = data_name_list
self.do_augment = do_augment self.do_augment = do_augment
self.random_annotations = random_annotations self.random_annotations = random_annotations
self.normalize_mean = normalize_mean self.normalize_mean = normalize_mean
@@ -48,12 +49,6 @@ class VCDADataset(Dataset):
# connect to database # connect to database
collection, _, _ = connect() collection, _, _ = connect()
self.collection = collection 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): def __len__(self):
return len(self.data_name_list) return len(self.data_name_list)