diff --git a/model/src/main.py b/model/src/main.py index 6a8b572..976b454 100644 --- a/model/src/main.py +++ b/model/src/main.py @@ -21,6 +21,7 @@ if __name__ == '__main__': model.eval() # setup dataset dataset = VCDADataset( + minio_client=minio_client, data_name_list=[ '02dbaf48d713e4e6d3a6b98fd2dc866e', ], diff --git a/model/src/utils/vcda_dataset.py b/model/src/utils/vcda_dataset.py index d5fc86a..edb5fc4 100644 --- a/model/src/utils/vcda_dataset.py +++ b/model/src/utils/vcda_dataset.py @@ -2,6 +2,7 @@ import random +from minio import Minio from PIL import Image from torch import Tensor from torch.utils.data import Dataset @@ -15,7 +16,7 @@ from torchvision.transforms.functional import ( to_tensor, ) -from shared.data_store import connect, get_image +from shared.data_store import get_image # resnet18 original normalization values RESNET_NORMALIZE_MEAN = [0.485, 0.456, 0.406] @@ -27,11 +28,13 @@ class VCDADataset(Dataset): def __init__( self, + minio_client: Minio, data_name_list: list[str], do_augment: bool = False, random_annotations: bool = False, ): super().__init__() + self.minio_client = minio_client self.data_name_list = data_name_list self.do_augment = do_augment self.random_annotations = random_annotations @@ -47,9 +50,6 @@ class VCDADataset(Dataset): contrast=8e-2, saturation=8e-2, ) - # connect to minio - minio_client = connect() - self.minio_client = minio_client def __len__(self): return len(self.data_name_list)