updated to use data_name_list
This commit is contained in:
@@ -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)
|
||||||
|
|||||||
Reference in New Issue
Block a user