removed unused packages

model/src/main.py

@@ -9,22 +9,16 @@ from models import VisualCommunicationModel
 from tqdm import tqdm
 from utils import DEVICE, VCDADataset, load_model
 
-from shared.datastore import Datastore
-from shared.docstore.classes import ModelData
 from shared.utils import setup_logging
 
 if __name__ == '__main__':
     # setup logging
     setup_logging()
-    # connect to minio
-    datastore = Datastore()
-    datastore.connect()
     # instantiate model
-    model: VisualCommunicationModel = load_model(client=datastore._client)
+    model: VisualCommunicationModel = load_model()
     model.eval()
     # setup dataset
     dataset = VCDADataset(
-        minio_client=datastore._client,
         data_name_list=[
             '02dbaf48d713e4e6d3a6b98fd2dc866e',
         ],
@@ -40,10 +34,10 @@ if __name__ == '__main__':
                 image = torch.unsqueeze(image, 0)  # add artificial batch dimension
                 image = image.to(DEVICE)
                 # make prediction
-                pred: ModelData = model(image)
+                pred: dict = model(image)
             except Exception:
                 print_exc()
                 continue
             else:
-                print(json.dumps(pred.model_dump(), indent=4))
+                print(json.dumps(pred, indent=4))
     logging.debug('finished')

model/src/utils/load_model.py

@@ -6,29 +6,28 @@ from pathlib import Path
 import torch
 
 from model.src.models import VisualCommunicationModel
-from shared.datastore import Datastore
+from shared.repositories import ModelRepository
 
 from .get_model_name import get_model_name
 
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 
-def load_model(
-    datastore: Datastore,
-) -> VisualCommunicationModel:
+def load_model() -> VisualCommunicationModel:
     """Instantiate model with weights loaded from latest model saved in
     MinIO."""
-    assert isinstance(datastore, Datastore)
     # instantiate model
     model = VisualCommunicationModel()
     # get model object name
     model_name_path = Path('model_name.txt')
     model_object_name = get_model_name(path=model_name_path)
     logging.info('using model: %s', model_object_name)
-    # load model from minio
-    model_checkpoint = datastore.get_model(
-        object_name=model_object_name,
-    )
+    # load model data
+    with ModelRepository() as repo:
+        model_data = repo.get_data(model_object_name)
+    if model_data is None:
+        raise FileNotFoundError(f'model {model_object_name} not found')
+    model_checkpoint = torch.load(model_data.buffer)
     model.load_state_dict(model_checkpoint)
     # clean memory
     model_checkpoint.clear()

model/src/utils/vcda_dataset.py

@@ -15,7 +15,7 @@ from torchvision.transforms.functional import (
     to_tensor,
 )
 
-from shared.datastore import Datastore
+from shared.repositories import ImageRepository
 
 # resnet18 original normalization values
 RESNET_NORMALIZE_MEAN = [0.485, 0.456, 0.406]
@@ -27,13 +27,11 @@ class VCDADataset(Dataset):
 
     def __init__(
         self,
-        datastore: Datastore,
         data_name_list: list[str],
         do_augment: bool = False,
         random_annotations: bool = False,
     ):
         super().__init__()
-        self.datastore = datastore
         self.data_name_list = data_name_list
         self.do_augment = do_augment
         self.random_annotations = random_annotations
@@ -55,11 +53,11 @@ class VCDADataset(Dataset):
 
     def __getitem__(self, idx):
         # get image from database
-        object_name = self.data_name_list[idx]
-        image = self.datastore.get_image(
-            object_name=object_name,
-        )
-        tensor = self.image_to_tensor(image)
+        image_name = self.data_name_list[idx]
+        with ImageRepository() as repo:
+            image_data = repo.get_data(image_name)
+        assert image_data is not None
+        tensor = self.image_to_tensor(image_data.image)
         if self.do_augment:
             tensor = self.augment(tensor)
         return tensor

model/train.py

@@ -18,7 +18,6 @@ from torch.optim.lr_scheduler import ExponentialLR
 from torch.utils.data import DataLoader
 
 from model.src.utils import VCDADataset, get_class
-from shared.datastore import Datastore
 
 
 def parse_arguments():
@@ -57,19 +56,16 @@ optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
 loss_fn = nn.CrossEntropyLoss()
 
 # create datasets and loaders
-datastore = Datastore()
-datastore.connect()
 with open('model/src/dataset/train.csv', encoding='utf-8') as fh:
     train_data_name_list = fh.read().split('\n')
 train_dataset = VCDADataset(
-    datastore=datastore,
     data_name_list=train_data_name_list,
 )
 train_loader = DataLoader(dataset=train_dataset, num_workers=args.loader_workers)
 
 with open('model/src/dataset/val.csv', encoding='utf-8') as fh:
     val_data_name_list = fh.read().split('\n')
-val_dataset = VCDADataset(datastore=datastore, data_name_list=val_data_name_list)
+val_dataset = VCDADataset(data_name_list=val_data_name_list)
 val_loader = DataLoader(dataset=val_dataset, num_workers=args.loader_workers)
 
 # create trainer and evaluator
@@ -141,7 +137,7 @@ to_save = {
 }
 checkpoint_handler = Checkpoint(
     to_save,
-    f"runs/checkpoints/{run_name}",
+    f'runs/checkpoints/{run_name}',
     n_saved=3,
     filename_prefix='best',
     score_function=lambda engine: -engine.state.metrics['loss'],
@@ -155,7 +151,7 @@ if args.checkpoint:
 
 # save model config
 os.makedirs('runs/configs/', exist_ok=True)
-with open(f"runs/configs/{run_name}.json", 'w', encoding='utf-8') as fh:
+with open(f'runs/configs/{run_name}.json', 'w', encoding='utf-8') as fh:
     json.dump(config, fh)
 
 # start training