datamint 1.9.2__py3-none-any.whl → 2.0.0__py3-none-any.whl

datamint/client_cmd_tools/datamint_config.py

@@ -1,46 +1,13 @@
 import argparse
 import logging
-import os
-import platform
 from datamint import configs
-from datamint.utils.logging_utils import load_cmdline_logging_config
+from datamint.utils.logging_utils import load_cmdline_logging_config, ConsoleWrapperHandler
 from rich.prompt import Prompt, Confirm
 from rich.console import Console
-from rich.theme import Theme
-
-# Create a custom theme that works well on both dark and blue backgrounds
-def _create_console_theme() -> Theme:
-    """Create a custom Rich theme optimized for cross-platform terminals."""
-    # Detect if we're likely on PowerShell (Windows + PowerShell)
-    is_powershell = (
-        platform.system() == "Windows" and 
-        os.environ.get("PSModulePath") is not None
-    )
-    
-    if is_powershell:
-        # PowerShell blue background - use high contrast colors
-        return Theme({
-            "warning": "bright_yellow",
-            "error": "bright_red on white",
-            "success": "bright_green",
-            "key": "bright_cyan",
-            "accent": "bright_cyan",
-            "title": "bold"
-        })
-    else:
-        # Linux/Unix terminals - standard colors
-        return Theme({
-            "warning": "yellow",
-            "error": "red",
-            "success": "green",
-            "key": "cyan",
-            "accent": "bright_blue",
-            "title": "bold"
-        })
-
-# Create console with custom theme
-console = Console(theme=_create_console_theme())
+
 _LOGGER = logging.getLogger(__name__)
+_USER_LOGGER = logging.getLogger('user_logger')
+console: Console
 
 
 def configure_default_url():
@@ -170,7 +137,9 @@ def interactive_mode():
 
 def main():
     """Main entry point for the configuration tool."""
+    global console
     load_cmdline_logging_config()
+    console = [h for h in _USER_LOGGER.handlers if isinstance(h, ConsoleWrapperHandler)][0].console
     parser = argparse.ArgumentParser(
         description='🔧 Datamint API Configuration Tool',
         epilog="""

datamint/client_cmd_tools/datamint_upload.py

@@ -1,6 +1,7 @@
 from datamint.exceptions import DatamintException
 import argparse
-from datamint.apihandler.api_handler import APIHandler
+# from datamint.apihandler.api_handler import APIHandler
+from datamint import Api
 import os
 from humanize import naturalsize
 import logging
@@ -12,8 +13,8 @@ from typing import Generator, Optional, Any
 from collections import defaultdict
 from datamint import __version__ as datamint_version
 from datamint import configs
-from datamint.client_cmd_tools.datamint_config import ask_api_key
-from datamint.utils.logging_utils import load_cmdline_logging_config
+from datamint.utils.logging_utils import load_cmdline_logging_config, ConsoleWrapperHandler
+from rich.console import Console
 import yaml
 from collections.abc import Iterable
 import pandas as pd
@@ -22,32 +23,42 @@ import pydicom.errors
 # Create two loggings: one for the user and one for the developer
 _LOGGER = logging.getLogger(__name__)
 _USER_LOGGER = logging.getLogger('user_logger')
+logging.getLogger('pydicom').setLevel(logging.ERROR)
+CONSOLE: Console
 
 MAX_RECURSION_LIMIT = 1000
 
+# Default extensions to exclude when --include-extensions is not specified
+DEFAULT_EXCLUDED_EXTENSIONS = [
+    '.txt', '.json', '.xml', '.docx', '.doc', '.pdf', '.xlsx', '.xls', '.csv', '.tsv',
+    '.log', '.ini', '.cfg', '.conf', '.yaml', '.yml', '.md', '.rst', '.html', '.htm',
+    '.exe', '.bat', '.sh', '.py', '.js', '.css',
+    '.sql', '.bak', '.tmp', '.temp', '.lock', '.DS_Store', '.gitignore'
+]
+
 
 def _get_minimal_distinguishing_paths(file_paths: list[str]) -> dict[str, str]:
     """
     Generate minimal distinguishing paths for files to avoid ambiguity when multiple files have the same name.
-    
+
     Args:
         file_paths: List of file paths
-        
+
     Returns:
         Dictionary mapping full path to minimal distinguishing path
     """
     if not file_paths:
         return {}
-    
+
     # Convert to Path objects and get absolute paths
     paths = [Path(fp).resolve() for fp in file_paths]
     result = {}
-    
+
     # Group files by basename
     basename_groups = defaultdict(list)
     for i, path in enumerate(paths):
         basename_groups[path.name].append((i, path))
-    
+
     for basename, path_list in basename_groups.items():
         if len(path_list) == 1:
             # Only one file with this name, use just the basename
@@ -56,7 +67,7 @@ def _get_minimal_distinguishing_paths(file_paths: list[str]) -> dict[str, str]:
         else:
             # Multiple files with same name, need to distinguish them
             path_parts_list = [path.parts for _, path in path_list]
-            
+
             # Find the minimum number of parent directories needed to distinguish
             max_depth_needed = 1
             for depth in range(1, max(len(parts) for parts in path_parts_list) + 1):
@@ -67,12 +78,12 @@ def _get_minimal_distinguishing_paths(file_paths: list[str]) -> dict[str, str]:
                         suffixes.append('/'.join(parts))
                     else:
                         suffixes.append('/'.join(parts[-depth:]))
-                
+
                 if len(set(suffixes)) == len(suffixes):
                     # All suffixes are unique at this depth
                     max_depth_needed = depth
                     break
-            
+
             # Apply the minimal distinguishing paths
             for (idx, path), parts in zip(path_list, path_parts_list):
                 if max_depth_needed >= len(parts):
@@ -80,10 +91,8 @@ def _get_minimal_distinguishing_paths(file_paths: list[str]) -> dict[str, str]:
                 else:
                     distinguishing_path = '/'.join(parts[-max_depth_needed:])
                 result[file_paths[idx]] = distinguishing_path
-    
-    return result
-
 
+    return result
 
 
 def _read_segmentation_names(segmentation_names_path: str | Path) -> dict:
@@ -257,6 +266,7 @@ def handle_api_key() -> str | None:
     If it does not exist, it asks the user to input it.
     Then, it asks the user if he wants to save the API key at a proper location in the machine
     """
+    from datamint.client_cmd_tools.datamint_config import ask_api_key
     api_key = configs.get_value(configs.APIKEY_KEY)
     if api_key is None:
         _USER_LOGGER.info("API key not found. Please provide it:")
@@ -524,6 +534,8 @@ def _parse_args() -> tuple[Any, list[str], Optional[list[dict]], Optional[list[s
     parser.add_argument('--channel', '--name', type=str, required=False,
                         help='Channel name (arbritary) to upload the resources to. \
                             Useful for organizing the resources in the platform.')
+    parser.add_argument('--project', type=str, required=False,
+                        help='Project name to add the uploaded resources to after successful upload.')
     parser.add_argument('--retain-pii', action='store_true', help='Do not anonymize DICOMs')
     parser.add_argument('--retain-attribute', type=_tuple_int_type, action='append',
                         default=[],
@@ -541,7 +553,8 @@ def _parse_args() -> tuple[Any, list[str], Optional[list[dict]], Optional[list[s
                         help='File extensions to be considered for uploading. Default: all file extensions.' +
                         ' Example: --include-extensions dcm jpg png')
     parser.add_argument('--exclude-extensions', type=str, nargs='+',
-                        help='File extensions to be excluded from uploading. Default: none.' +
+                        help='File extensions to be excluded from uploading. ' +
+                        'Default: common non-medical file extensions (.txt, .json, .xml, .docx, etc.) when --include-extensions is not specified.' +
                         ' Example: --exclude-extensions txt csv'
                         )
     parser.add_argument('--segmentation_path', type=_is_valid_path_argparse, metavar="FILE",
@@ -581,7 +594,6 @@ def _parse_args() -> tuple[Any, list[str], Optional[list[dict]], Optional[list[s
 
     if args.verbose:
         # Get the console handler and set to debug
-        print(logging.getLogger().handlers)
         logging.getLogger().handlers[0].setLevel(logging.DEBUG)
         logging.getLogger('datamint').setLevel(logging.DEBUG)
         _LOGGER.setLevel(logging.DEBUG)
@@ -594,6 +606,11 @@ def _parse_args() -> tuple[Any, list[str], Optional[list[dict]], Optional[list[s
     if args.include_extensions is not None and args.exclude_extensions is not None:
         raise ValueError("--include-extensions and --exclude-extensions are mutually exclusive.")
 
+    # Apply default excluded extensions if neither include nor exclude extensions are specified
+    if args.include_extensions is None and args.exclude_extensions is None:
+        args.exclude_extensions = DEFAULT_EXCLUDED_EXTENSIONS
+        _LOGGER.debug(f"Applied default excluded extensions: {args.exclude_extensions}")
+
     try:
         if os.path.isfile(args.path):
             file_path = [args.path]
@@ -720,24 +737,26 @@ def print_results_summary(files_path: list[str],
     failure_files = [f for f, r in zip(files_path, results) if isinstance(r, Exception)]
     # Get distinguishing paths for better error reporting
     distinguishing_paths = _get_minimal_distinguishing_paths(files_path)
-    
+
     _USER_LOGGER.info(f"\nUpload summary:")
     _USER_LOGGER.info(f"\tTotal files: {len(files_path)}")
     _USER_LOGGER.info(f"\tSuccessful uploads: {len(files_path) - len(failure_files)}")
     if len(failure_files) > 0:
-        _USER_LOGGER.info(f"\t❌ Failed uploads: {len(failure_files)}")
+        _USER_LOGGER.warning(f"\tFailed uploads: {len(failure_files)}")
         _USER_LOGGER.warning(f"\tFailed files: {[distinguishing_paths[f] for f in failure_files]}")
         _USER_LOGGER.warning(f"\nFailures:")
         for f, r in zip(files_path, results):
             if isinstance(r, Exception):
                 _USER_LOGGER.warning(f"\t{distinguishing_paths[f]}: {r}")
     else:
-        _USER_LOGGER.info(f'✅ All uploads successful!')
+        CONSOLE.print(f'✅ All uploads successful!', style='success')
     return len(failure_files)
 
 
 def main():
+    global CONSOLE
     load_cmdline_logging_config()
+    CONSOLE = [h for h in _USER_LOGGER.handlers if isinstance(h, ConsoleWrapperHandler)][0].console
 
     try:
         args, files_path, segfiles, metadata_files = _parse_args()
@@ -745,48 +764,55 @@ def main():
         _USER_LOGGER.error(f'Error validating arguments. {e}')
         sys.exit(1)
 
-    print_input_summary(files_path,
-                        args=args,
-                        segfiles=segfiles,
-                        metadata_files=metadata_files,
-                        include_extensions=args.include_extensions)
+    try:
+        print_input_summary(files_path,
+                            args=args,
+                            segfiles=segfiles,
+                            metadata_files=metadata_files,
+                            include_extensions=args.include_extensions)
 
-    if not args.yes:
-        confirmation = input("Do you want to proceed with the upload? (y/n): ")
-        if confirmation.lower() != "y":
-            _USER_LOGGER.info("Upload cancelled.")
-            return
-    #######################################
+        if not args.yes:
+            confirmation = input("Do you want to proceed with the upload? (y/n): ")
+            if confirmation.lower() != "y":
+                _USER_LOGGER.info("Upload cancelled.")
+                return
+        #######################################
 
-    has_a_dicom_file = any(is_dicom(f) for f in files_path)
+        has_a_dicom_file = any(is_dicom(f) for f in files_path)
 
-    try:
-        api_handler = APIHandler(check_connection=True)
-    except DatamintException as e:
-        _USER_LOGGER.error(f'❌ Connection failed: {e}')
-        return
-    try:
-        results = api_handler.upload_resources(channel=args.channel,
-                                               files_path=files_path,
-                                               tags=args.tag,
-                                               on_error='skip',
-                                               anonymize=args.retain_pii == False and has_a_dicom_file,
-                                               anonymize_retain_codes=args.retain_attribute,
-                                               mung_filename=args.mungfilename,
-                                               publish=args.publish,
-                                               segmentation_files=segfiles,
-                                               transpose_segmentation=args.transpose_segmentation,
-                                               assemble_dicoms=True,
-                                               metadata=metadata_files
-                                               )
-    except pydicom.errors.InvalidDicomError as e:
-        _USER_LOGGER.error(f'❌ Invalid DICOM file: {e}')
-        return
-    _USER_LOGGER.info('Upload finished!')
-    _LOGGER.debug(f"Number of results: {len(results)}")
-
-    num_failures = print_results_summary(files_path, results)
-    if num_failures > 0:
+        try:
+            api = Api(check_connection=True)
+        except DatamintException as e:
+            _USER_LOGGER.error(f'❌ Connection failed: {e}')
+            return
+        try:
+            print('>>>', segfiles)
+            results = api.resources.upload_resources(channel=args.channel,
+                                                     files_path=files_path,
+                                                     tags=args.tag,
+                                                     on_error='skip',
+                                                     anonymize=args.retain_pii == False and has_a_dicom_file,
+                                                     anonymize_retain_codes=args.retain_attribute,
+                                                     mung_filename=args.mungfilename,
+                                                     publish=args.publish,
+                                                     publish_to=args.project,
+                                                     segmentation_files=segfiles,
+                                                     transpose_segmentation=args.transpose_segmentation,
+                                                     assemble_dicoms=True,
+                                                     metadata=metadata_files,
+                                                     progress_bar=True
+                                                     )
+        except pydicom.errors.InvalidDicomError as e:
+            _USER_LOGGER.error(f'❌ Invalid DICOM file: {e}')
+            return
+        _USER_LOGGER.info('Upload finished!')
+        _LOGGER.debug(f"Number of results: {len(results)}")
+
+        num_failures = print_results_summary(files_path, results)
+        if num_failures > 0:
+            sys.exit(1)
+    except KeyboardInterrupt:
+        CONSOLE.print("\nUpload cancelled by user.", style='warning')
         sys.exit(1)
 
 

datamint/dataset/base_dataset.py

@@ -13,14 +13,16 @@ from datamint import configs
 from torch.utils.data import DataLoader
 import torch
 from torch import Tensor
-from datamint.apihandler.base_api_handler import DatamintException
+from datamint.exceptions import DatamintException
 from medimgkit.dicom_utils import is_dicom
 from medimgkit.readers import read_array_normalized
-from medimgkit.format_detection import guess_extension
+from medimgkit.format_detection import guess_extension, guess_typez
+from medimgkit.nifti_utils import NIFTI_MIMES, get_nifti_shape
 from datetime import datetime
 from pathlib import Path
-from datamint.dataset.annotation import Annotation
+from datamint.entities import Annotation, DatasetInfo
 import cv2
+from datamint.entities import Resource
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -174,23 +176,12 @@ class DatamintBaseDataset:
 
     def _setup_api_handler(self, server_url: Optional[str], api_key: Optional[str], auto_update: bool) -> None:
         """Setup API handler and validate connection."""
-        from datamint.apihandler.api_handler import APIHandler
-
-        self.api_handler = APIHandler(
-            root_url=server_url,
+        from datamint import Api
+        self.api = Api(
+            server_url=server_url,
             api_key=api_key,
-            check_connection=auto_update
+            check_connection=self.auto_update
         )
-        self.server_url = self.api_handler.root_url
-        self.api_key = self.api_handler.api_key
-
-        if self.api_key is None:
-            _LOGGER.warning(
-                "API key not provided. If you want to download data, please provide an API key, "
-                f"either by passing it as an argument, "
-                f"setting environment variable {configs.ENV_VARS[configs.APIKEY_KEY]} or "
-                "using datamint-config command line tool."
-            )
 
     def _setup_directories(self, root: str | None) -> None:
         """Setup root and dataset directories."""
@@ -242,7 +233,7 @@ class DatamintBaseDataset:
         if not os.path.isfile(metadata_path):
             # get the server info
             self.project_info = self.get_info()
-            self.metainfo = self._get_datasetinfo().copy()
+            self.metainfo = self._get_datasetinfo().asdict().copy()
             self.metainfo['updated_at'] = None
             self.metainfo['resources'] = []
             self.metainfo['all_annotations'] = self.all_annotations
@@ -412,19 +403,33 @@ class DatamintBaseDataset:
     @staticmethod
     def read_number_of_frames(filepath: str) -> int:
         """Read the number of frames in a file."""
-        if is_dicom(filepath):
+
+        mimetypes, ext = guess_typez(filepath)
+        mimetype = mimetypes[0]
+        if mimetype is None:
+            raise ValueError(f"Could not determine MIME type for file: {filepath}")
+
+        if mimetype == 'application/dicom':
             ds = pydicom.dcmread(filepath)
             return getattr(ds, 'NumberOfFrames', 1)
-        elif filepath.lower().endswith(('.mp4', '.avi')):
+        elif mimetype.startswith('video/'):
             cap = cv2.VideoCapture(filepath)
             try:
                 return int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
             finally:
                 cap.release()
-        elif filepath.lower().endswith(('.png', '.jpg', '.jpeg')):
+        elif mimetype in ('image/png', 'image/jpeg', 'image/jpg', 'image/bmp', 'image/tiff'):
             return 1
+        elif mimetype in NIFTI_MIMES:
+            shape = get_nifti_shape(filepath)
+            if len(shape) == 3:
+                return shape[-1]
+            elif len(shape) > 3:
+                return shape[3]
+            else:
+                return 1
         else:
-            raise ValueError(f"Unsupported file type: {filepath}")
+            raise ValueError(f"Unsupported file type '{mimetype}' for file {filepath}")
 
     def get_resources_ids(self) -> list[str]:
         """Get list of resource IDs."""
@@ -526,18 +531,18 @@ class DatamintBaseDataset:
         if missing_files:
             raise DatamintDatasetException(f"Image files not found: {missing_files}")
 
-    def _get_datasetinfo(self) -> dict:
+    def _get_datasetinfo(self) -> DatasetInfo:
         """Get dataset information from API."""
         if self._server_dataset_info is not None:
             return self._server_dataset_info
-        all_datasets = self.api_handler.get_datasets()
+        all_datasets = self.api._datasetsinfo.get_all()
 
         for dataset in all_datasets:
-            if dataset['id'] == self.dataset_id:
+            if dataset.id == self.dataset_id:
                 self._server_dataset_info = dataset
                 return dataset
 
-        available_datasets = [(d['name'], d['id']) for d in all_datasets]
+        available_datasets = [(d.name, d.id) for d in all_datasets]
         raise DatamintDatasetException(
             f"Dataset with id '{self.dataset_id}' not found. "
             f"Available datasets: {available_datasets}"
@@ -547,7 +552,7 @@ class DatamintBaseDataset:
         """Get project information from API."""
         if hasattr(self, 'project_info') and self.project_info is not None:
             return self.project_info
-        project = self.api_handler.get_project_by_name(self.project_name)
+        project = self.api.projects.get_by_name(self.project_name).asdict()
         if 'error' in project:
             available_projects = project['all_projects']
             raise DatamintDatasetException(
@@ -592,31 +597,10 @@ class DatamintBaseDataset:
         lines = [head] + [" " * 4 + line for line in body]
         return "\n".join(lines)
 
-    def download_project(self) -> None:
-        """Download project data from API."""
-
-        dataset_info = self._get_datasetinfo()
-        self.dataset_id = dataset_info['id']
-        self.last_updaded_at = dataset_info['updated_at']
-
-        self.api_handler.download_project(
-            self.project_info['id'],
-            self.dataset_zippath,
-            all_annotations=self.all_annotations,
-            include_unannotated=self.include_unannotated
-        )
-
-        _LOGGER.debug("Downloaded dataset")
-
-        if os.path.getsize(self.dataset_zippath) == 0:
-            raise DatamintDatasetException("Download failed.")
-
-        self._extract_and_update_metadata()
-
     def _get_dataset_id(self) -> str:
         if self.dataset_id is None:
             dataset_info = self._get_datasetinfo()
-            self.dataset_id = dataset_info['id']
+            self.dataset_id = dataset_info.id
         return self.dataset_id
 
     def _extract_and_update_metadata(self) -> None:
@@ -638,7 +622,7 @@ class DatamintBaseDataset:
 
         # Save updated metadata
         with open(datasetjson_path, 'w') as file:
-            json.dump(self.metainfo, file, default=lambda o: o.to_dict() if hasattr(o, 'to_dict') else o)
+            json.dump(self.metainfo, file, default=lambda o: o.asdict() if hasattr(o, 'asdict') else o)
 
         self.images_metainfo = self.metainfo['resources']
         # self._convert_metainfo_to_clsobj()
@@ -646,19 +630,19 @@ class DatamintBaseDataset:
     def _update_metadata_timestamps(self) -> None:
         """Update metadata with correct timestamps."""
         if 'updated_at' not in self.metainfo:
-            self.metainfo['updated_at'] = self.last_updaded_at
+            self.metainfo['updated_at'] = self.last_updated_at
         else:
             try:
                 local_time = datetime.fromisoformat(self.metainfo['updated_at'])
-                server_time = datetime.fromisoformat(self.last_updaded_at)
+                server_time = datetime.fromisoformat(self.last_updated_at)
 
                 if local_time < server_time:
                     _LOGGER.warning(
                         f"Inconsistent updated_at dates detected "
-                        f"({self.metainfo['updated_at']} < {self.last_updaded_at}). "
-                        f"Fixing it to {self.last_updaded_at}"
+                        f"({self.metainfo['updated_at']} < {self.last_updated_at}). "
+                        f"Fixing it to {self.last_updated_at}"
                     )
-                    self.metainfo['updated_at'] = self.last_updaded_at
+                    self.metainfo['updated_at'] = self.last_updated_at
             except Exception as e:
                 _LOGGER.warning(f"Failed to parse updated_at date: {e}")
 
@@ -690,6 +674,9 @@ class DatamintBaseDataset:
             img = (img - min_val) / (img.max() - min_val) * 255
             img = img.astype(np.uint8)
 
+        if not img.flags.writeable:
+            img = img.copy()
+
         img_tensor = torch.from_numpy(img).contiguous()
 
         if isinstance(img_tensor, torch.ByteTensor):
@@ -829,7 +816,7 @@ class DatamintBaseDataset:
 
         try:
             external_metadata_info = self._get_datasetinfo()
-            server_updated_at = external_metadata_info['updated_at']
+            server_updated_at = external_metadata_info.updated_at
         except Exception as e:
             _LOGGER.warning(f"Failed to check for updates in {self.project_name}: {e}")
             return
@@ -856,20 +843,21 @@ class DatamintBaseDataset:
             _LOGGER.info('Local version is up to date with the latest version.')
 
     def _fetch_new_resources(self,
-                             all_uptodate_resources: list[dict]) -> list[dict]:
+                             all_uptodate_resources: list[Resource]) -> list[dict]:
         local_resources = self.images_metainfo
         local_resources_ids = [res['id'] for res in local_resources]
         new_resources = []
         for resource in all_uptodate_resources:
+            resource = resource.asdict()
             if resource['id'] not in local_resources_ids:
                 resource['file'] = str(self._get_resource_file_path(resource))
                 resource['annotations'] = []
                 new_resources.append(resource)
         return new_resources
 
-    def _fetch_deleted_resources(self, all_uptodate_resources: list[dict]) -> list[dict]:
+    def _fetch_deleted_resources(self, all_uptodate_resources: list[Resource]) -> list[dict]:
         local_resources = self.images_metainfo
-        all_uptodate_resources_ids = [res['id'] for res in all_uptodate_resources]
+        all_uptodate_resources_ids = [res.id for res in all_uptodate_resources]
         deleted_resources = []
         for resource in local_resources:
             try:
@@ -888,7 +876,7 @@ class DatamintBaseDataset:
         # server_updated_at = external_metadata_info['updated_at']
 
         ### RESOURCES ###
-        all_uptodate_resources = self.api_handler.get_project_resources(self.get_info()['id'])
+        all_uptodate_resources = self.api.projects.get_project_resources(self.get_info()['id'])
         new_resources = self._fetch_new_resources(all_uptodate_resources)
         deleted_resources = self._fetch_deleted_resources(all_uptodate_resources)
 
@@ -898,9 +886,9 @@ class DatamintBaseDataset:
             new_resources_path = [Path(self.dataset_dir) / r['file'] for r in new_resources]
             new_resources_ids = [r['id'] for r in new_resources]
             _LOGGER.info(f"Downloading {len(new_resources)} new resources...")
-            new_res_paths = self.api_handler.download_multiple_resources(new_resources_ids,
-                                                                         save_path=new_resources_path,
-                                                                         add_extension=True)
+            new_res_paths = self.api.resources.download_multiple_resources(new_resources_ids,
+                                                                           save_path=new_resources_path,
+                                                                           add_extension=True)
             for new_rpath, r in zip(new_res_paths, new_resources):
                 r['file'] = str(Path(new_rpath).relative_to(self.dataset_dir))
             _LOGGER.info(f"Downloaded {len(new_resources)} new resources.")
@@ -910,16 +898,17 @@ class DatamintBaseDataset:
         ################
 
         ### ANNOTATIONS ###
-        all_annotations = self.api_handler.get_annotations(worklist_id=self.project_info['worklist_id'],
-                                                           status='published' if self.all_annotations else None)
+        all_annotations = self.api.annotations.get_list(worklist_id=self.project_info['worklist_id'],
+                                                        status='published' if self.all_annotations else None)
+
         # group annotations by resource ID
-        annotations_by_resource = {}
+        annotations_by_resource: dict[str, list[Annotation]] = {}
         for ann in all_annotations:
             # add the local filepath
             filepath = self._get_annotation_file_path(ann)
             if filepath is not None:
-                ann['file'] = str(filepath)
-            resource_id = ann['resource_id']
+                ann.file = str(filepath)
+            resource_id = ann.resource_id
             if resource_id not in annotations_by_resource:
                 annotations_by_resource[resource_id] = []
             annotations_by_resource[resource_id].append(ann)
@@ -937,11 +926,11 @@ class DatamintBaseDataset:
             # check if segmentation annotations need to be downloaded
             # Also check if annotations need to be deleted
             old_ann_ids = set([ann.id for ann in old_resource_annotations if hasattr(ann, 'id')])
-            new_ann_ids = set([ann['id'] for ann in new_resource_annotations])
+            new_ann_ids = set([ann.id for ann in new_resource_annotations])
 
             # Find annotations to add, update, or remove
             annotations_to_add = [ann for ann in new_resource_annotations
-                                  if ann['id'] not in old_ann_ids]
+                                  if ann.id not in old_ann_ids]
             annotations_to_remove = [ann for ann in old_resource_annotations
                                      if getattr(ann, 'id', 'NA') not in new_ann_ids]
 
@@ -970,22 +959,23 @@ class DatamintBaseDataset:
                     _LOGGER.error(f"Error deleting annotation file {filepath}: {e}")
 
             # Update resource annotations list - convert to Annotation objects
-            resource['annotations'] = [Annotation.from_dict(ann) for ann in new_resource_annotations]
+            # resource['annotations'] = [Annotation.from_dict(ann) for ann in new_resource_annotations]
+            resource['annotations'] = new_resource_annotations
 
         # Batch download all segmentation files
         if segmentations_to_download:
             _LOGGER.info(f"Downloading {len(segmentations_to_download)} segmentation files...")
-            self.api_handler.download_multiple_segmentations(segmentations_to_download, segmentation_paths)
+            self.api.annotations.download_multiple_files(segmentations_to_download, segmentation_paths)
             _LOGGER.info(f"Downloaded {len(segmentations_to_download)} segmentation files.")
 
         ###################
         # update metadata
-        self.metainfo['updated_at'] = self._get_datasetinfo()['updated_at']
+        self.metainfo['updated_at'] = self._get_datasetinfo().updated_at
         self.metainfo['all_annotations'] = self.all_annotations
         # save updated metadata
         datasetjson_path = os.path.join(self.dataset_dir, 'dataset.json')
         with open(datasetjson_path, 'w') as file:
-            json.dump(self.metainfo, file, default=lambda o: o.to_dict() if hasattr(o, 'to_dict') else o)
+            json.dump(self.metainfo, file, default=lambda o: o.asdict() if hasattr(o, 'asdict') else o)
 
     def _get_resource_file_path(self, resource: dict) -> Path:
         """Get the local file path for a resource."""