From aa3c1acb5412437e16861bedb4ad6b7cf3241256 Mon Sep 17 00:00:00 2001 From: Lucashsmello Date: Fri, 1 Aug 2025 16:30:58 -0300 Subject: [PATCH 1/5] Add Annotation class and refactor dataset handling for incremental update - Introduced a new `Annotation` class to encapsulate annotation data and provide methods for serialization and deserialization. - Updated `DatamintBaseDataset` to utilize the new `Annotation` class, converting annotations from dictionaries to `Annotation` objects. - Refactored methods to handle annotations, including filtering and loading segmentations, to work with the new class structure. - Enhanced dataset metadata handling to support incremental updates for new and deleted resources and annotations. - Improved error handling and logging for better traceability during dataset operations. --- datamint/apihandler/annotation_api_handler.py | 59 +- datamint/apihandler/base_api_handler.py | 56 +- datamint/apihandler/root_api_handler.py | 140 +++-- datamint/dataset/annotation.py | 221 ++++++++ datamint/dataset/base_dataset.py | 506 ++++++++++++------ datamint/dataset/dataset.py | 27 +- 6 files changed, 771 insertions(+), 238 deletions(-) create mode 100644 datamint/dataset/annotation.py diff --git a/datamint/apihandler/annotation_api_handler.py b/datamint/apihandler/annotation_api_handler.py index 8759ab3e..ec182bb3 100644 --- a/datamint/apihandler/annotation_api_handler.py +++ b/datamint/apihandler/annotation_api_handler.py @@ -13,6 +13,8 @@ from .dto.annotation_dto import CreateAnnotationDto, LineGeometry, BoxGeometry, CoordinateSystem, AnnotationType import pydicom import json +from deprecated import deprecated +from pathlib import Path _LOGGER = logging.getLogger(__name__) _USER_LOGGER = logging.getLogger('user_logger') @@ -267,8 +269,9 @@ async def _upload_volume_segmentation_async(self, raise NotImplementedError("`name=string` is not supported yet for volume segmentation.") if isinstance(name, dict): if any(isinstance(k, tuple) for k in name.keys()): - raise NotImplementedError("For volume segmentations, `name` must be a dictionary with integer keys only.") - + raise NotImplementedError( + "For volume segmentations, `name` must be a dictionary with integer keys only.") + # Prepare file for upload if isinstance(file_path, str): if file_path.endswith('.nii') or file_path.endswith('.nii.gz'): @@ -1098,6 +1101,29 @@ def delete_annotation(self, annotation_id: str | dict): resp = self._run_request(request_params) self._check_errors_response_json(resp) + def get_annotation_by_id(self, annotation_id: str) -> dict: + """ + Get an annotation by its unique id. + + Args: + annotation_id (str): The annotation unique id. + + Returns: + dict: The annotation information. + """ + request_params = { + 'method': 'GET', + 'url': f'{self.root_url}/annotations/{annotation_id}', + } + + try: + resp = self._run_request(request_params) + return resp.json() + except HTTPError as e: + _LOGGER.error(f"Error getting annotation by id {annotation_id}: {e}") + raise + + @deprecated(reason="Use download_segmentation_file instead") def get_segmentation_file(self, resource_id: str, annotation_id: str) -> bytes: request_params = { 'method': 'GET', @@ -1107,6 +1133,35 @@ def get_segmentation_file(self, resource_id: str, annotation_id: str) -> bytes: resp = self._run_request(request_params) return resp.content + def download_segmentation_file(self, annotation: str | dict, fpath_out: str | Path | None) -> bytes: + """ + Download the segmentation file for a given resource and annotation. + + Args: + annotation (str | dict): The annotation unique id or an annotation object. + fpath_out (str | None): (Optional) The file path to save the downloaded segmentation file. + + Returns: + bytes: The content of the downloaded segmentation file in bytes format. + """ + if isinstance(annotation, dict): + annotation_id = annotation['id'] + resource_id = annotation['resource_id'] + else: + annotation_id = annotation + resource_id = self.get_annotation_by_id(annotation_id)['resource_id'] + + request_params = { + 'method': 'GET', + 'url': f'{self.root_url}/annotations/{resource_id}/annotations/{annotation_id}/file', + } + + resp = self._run_request(request_params) + if fpath_out is not None: + with open(str(fpath_out), 'wb') as f: + f.write(resp.content) + return resp.content + def set_annotation_status(self, project_id: str, resource_id: str, diff --git a/datamint/apihandler/base_api_handler.py b/datamint/apihandler/base_api_handler.py index b7820e74..73eb8398 100644 --- a/datamint/apihandler/base_api_handler.py +++ b/datamint/apihandler/base_api_handler.py @@ -85,7 +85,7 @@ def __init__(self, msg = f"API key not provided! Use the environment variable " + \ f"{BaseAPIHandler.DATAMINT_API_VENV_NAME} or pass it as an argument." raise DatamintException(msg) - self.semaphore = asyncio.Semaphore(10) # Limit to 10 parallel requests + self.semaphore = asyncio.Semaphore(20) if check_connection: self.check_connection() @@ -157,30 +157,34 @@ def _generate_curl_command(self, request_args: dict) -> str: async def _run_request_async(self, request_args: dict, session: aiohttp.ClientSession | None = None, - data_to_get: str = 'json'): + data_to_get: Literal['json', 'text', 'content'] = 'json'): if session is None: async with aiohttp.ClientSession() as s: - return await self._run_request_async(request_args, s) - try: - _LOGGER.debug(f"Running request to {request_args['url']}") - _LOGGER.debug(f'Equivalent curl command: "{self._generate_curl_command(request_args)}"') - except Exception as e: - _LOGGER.debug(f"Error generating curl command: {e}") - - # add apikey to the headers - if 'headers' not in request_args: - request_args['headers'] = {} - - request_args['headers']['apikey'] = self.api_key - - async with session.request(**request_args) as response: - self._check_errors_response(response, request_args) - if data_to_get == 'json': - return await response.json() - elif data_to_get == 'text': - return await response.text() - else: - raise ValueError("data_to_get must be either 'json' or 'text'") + return await self._run_request_async(request_args, s, data_to_get) + + async with self.semaphore: + try: + _LOGGER.debug(f"Running request to {request_args['url']}") + _LOGGER.debug(f'Equivalent curl command: "{self._generate_curl_command(request_args)}"') + except Exception as e: + _LOGGER.debug(f"Error generating curl command: {e}") + + # add apikey to the headers + if 'headers' not in request_args: + request_args['headers'] = {} + + request_args['headers']['apikey'] = self.api_key + + async with session.request(**request_args) as response: + self._check_errors_response(response, request_args) + if data_to_get == 'json': + return await response.json() + elif data_to_get == 'text': + return await response.text() + elif data_to_get == 'content': + return await response.read() + else: + raise ValueError("data_to_get must be either 'json' or 'text'") def _check_errors_response(self, response, @@ -237,9 +241,9 @@ def _get_endpoint_url(self, endpoint: str) -> str: return f'{self.root_url}/{endpoint}' def _run_pagination_request(self, - request_params: Dict, - return_field: Optional[Union[str, List]] = None - ) -> Generator[Dict, None, None]: + request_params: dict, + return_field: str | list | None = None + ) -> Generator[dict | list, None, None]: offset = 0 params = request_params.get('params', {}) while True: diff --git a/datamint/apihandler/root_api_handler.py b/datamint/apihandler/root_api_handler.py index bc2dd95d..f5c76f12 100644 --- a/datamint/apihandler/root_api_handler.py +++ b/datamint/apihandler/root_api_handler.py @@ -219,36 +219,35 @@ async def _upload_resources_async(self, async with aiohttp.ClientSession() as session: async def __upload_single_resource(file_path, segfiles: dict[str, list | dict], metadata_file: str | dict | None): - async with self.semaphore: - rid = await self._upload_single_resource_async( - file_path=file_path, - mimetype=mimetype, - anonymize=anonymize, - anonymize_retain_codes=anonymize_retain_codes, - tags=tags, - session=session, - mung_filename=mung_filename, - channel=channel, - modality=modality, - publish=publish, - metadata_file=metadata_file, - ) - if segfiles is not None: - fpaths = segfiles['files'] - names = segfiles.get('names', _infinite_gen(None)) - if isinstance(names, dict): - names = _infinite_gen(names) - frame_indices = segfiles.get('frame_index', _infinite_gen(None)) - for f, name, frame_index in tqdm(zip(fpaths, names, frame_indices), - desc=f"Uploading segmentations for {file_path}", - total=len(fpaths)): - if f is not None: - await self._upload_segmentations_async(rid, - file_path=f, - name=name, - frame_index=frame_index, - transpose_segmentation=transpose_segmentation) - return rid + rid = await self._upload_single_resource_async( + file_path=file_path, + mimetype=mimetype, + anonymize=anonymize, + anonymize_retain_codes=anonymize_retain_codes, + tags=tags, + session=session, + mung_filename=mung_filename, + channel=channel, + modality=modality, + publish=publish, + metadata_file=metadata_file, + ) + if segfiles is not None: + fpaths = segfiles['files'] + names = segfiles.get('names', _infinite_gen(None)) + if isinstance(names, dict): + names = _infinite_gen(names) + frame_indices = segfiles.get('frame_index', _infinite_gen(None)) + for f, name, frame_index in tqdm(zip(fpaths, names, frame_indices), + desc=f"Uploading segmentations for {file_path}", + total=len(fpaths)): + if f is not None: + await self._upload_segmentations_async(rid, + file_path=f, + name=name, + frame_index=frame_index, + transpose_segmentation=transpose_segmentation) + return rid tasks = [__upload_single_resource(f, segfiles, metadata_file) for f, segfiles, metadata_file in zip(files_path, segmentation_files, metadata_files)] @@ -830,6 +829,62 @@ def set_resource_tags(self, def _has_status_code(e, status_code: int) -> bool: return hasattr(e, 'response') and (e.response is not None) and e.response.status_code == status_code + async def _async_download_file(self, + resource_id: str, + save_path: str, + session: aiohttp.ClientSession | None = None, + progress_bar: tqdm | None = None): + """ + Asynchronously download a file from the server. + + Args: + resource_id (str): The resource unique id. + save_path (str): The path to save the file. + session (aiohttp.ClientSession): The aiohttp session to use for the request. + progress_bar (tqdm | None): Optional progress bar to update after download completion. + """ + url = f"{self._get_endpoint_url(RootAPIHandler.ENDPOINT_RESOURCES)}/{resource_id}/file" + request_params = { + 'method': 'GET', + 'headers': {'accept': 'application/octet-stream'}, + 'url': url + } + try: + data_bytes = await self._run_request_async(request_params, session, 'content') + with open(save_path, 'wb') as f: + f.write(data_bytes) + if progress_bar: + progress_bar.update(1) + except ResourceNotFoundError as e: + e.set_params('resource', {'resource_id': resource_id}) + raise e + + def download_multiple_resources(self, + resource_ids: list[str], + save_path: list[str] | str + ) -> None: + """ + Download multiple resources and save them to the specified paths. + + Args: + resource_ids (list[str]): A list of resource unique ids. + save_path (list[str] | str): A list of paths to save the files or a directory path. + """ + async def _download_all_async(): + async with aiohttp.ClientSession() as session: + tasks = [ + self._async_download_file(resource_id, save_path=path, session=session, progress_bar=progress_bar) + for resource_id, path in zip(resource_ids, save_path) + ] + await asyncio.gather(*tasks) + + if isinstance(save_path, str): + save_path = [os.path.join(save_path, r) for r in resource_ids] + + with tqdm(total=len(resource_ids), desc="Downloading resources", unit="file") as progress_bar: + loop = asyncio.get_event_loop() + loop.run_until_complete(_download_all_async()) + def download_resource_file(self, resource_id: str, save_path: Optional[str] = None, @@ -988,6 +1043,7 @@ def get_datasets(self) -> list[dict]: response = self._run_request(request_params) return response.json()['data'] + @deprecated(version='1.7') def get_datasetsinfo_by_name(self, dataset_name: str) -> list[dict]: request_params = { 'method': 'GET', @@ -1082,6 +1138,30 @@ def get_projects(self) -> list[dict]: } return self._run_request(request_params).json()['data'] + def get_project_resources(self, project_id: str) -> list[dict]: + """ + Get the resources of a project by its id. + + Args: + project_id (str): The project id. + + Returns: + list[dict]: The list of resources in the project. + + Raises: + ResourceNotFoundError: If the project does not exists. + """ + request_params = { + 'method': 'GET', + 'url': f'{self.root_url}/projects/{project_id}/resources' + } + try: + return self._run_request(request_params).json() + except HTTPError as e: + if e.response is not None and e.response.status_code == 500: + raise ResourceNotFoundError('project', {'project_id': project_id}) + raise e + def create_project(self, name: str, description: str, diff --git a/datamint/dataset/annotation.py b/datamint/dataset/annotation.py new file mode 100644 index 00000000..11c91537 --- /dev/null +++ b/datamint/dataset/annotation.py @@ -0,0 +1,221 @@ +from __future__ import annotations +from dataclasses import dataclass, field, asdict +from typing import Optional, Any, TYPE_CHECKING +from datetime import datetime +from pathlib import Path +import logging +import numpy as np +from PIL import Image +import json + +# if TYPE_CHECKING: +# from datamint.apihandler.annotation_api_handler import AnnotationAPIHandler + +_LOGGER = logging.getLogger(__name__) + + +# Map API field names to class attributes +_FIELD_MAPPING = { + 'type':'annotation_type', + 'name': 'identifier', + 'added_by': 'created_by', + 'index': 'frame_index', +} + +@dataclass +class Annotation: + """ + Class representing an annotation from the Datamint API. + + This class stores annotation data and provides methods for loading + and saving annotations through the API handler. + + Args: + id: Unique identifier for the annotation + identifier: The annotation identifier/label name + scope: Whether annotation applies to 'frame' or 'image' + annotation_type: Type of annotation ('segmentation', 'label', 'category', etc.) + resource_id: ID of the resource this annotation belongs to + annotation_worklist_id: ID of the annotation worklist + created_by: Email of the user who created the annotation + status: Status of the annotation ('published', 'new', etc.) + frame_index: Frame index for frame-scoped annotations + text_value: Text value for category annotations + numeric_value: Numeric value for numeric annotations + units: Units for numeric annotations + geometry: Geometry data for geometric annotations + created_at: When the annotation was created + approved_at: When the annotation was approved + approved_by: Who approved the annotation + associated_file: Path to associated file (for segmentations) + deleted: Whether the annotation is deleted + deleted_at: When the annotation was deleted + deleted_by: Who deleted the annotation + created_by_model: Model ID if created by AI + old_geometry: Previous geometry data + set_name: Set name for grouped annotations + resource_filename: Filename of the associated resource + resource_modality: Modality of the associated resource + annotation_worklist_name: Name of the annotation worklist + user_info: Information about the user who created the annotation + values: Additional values + """ + + id: str + identifier: str + scope: str + annotation_type: str + resource_id: str + created_by: str + annotation_worklist_id: Optional[str] = None + status: Optional[str] = None + frame_index: Optional[int] = None + text_value: Optional[str] = None + numeric_value: Optional[float] = None + units: Optional[str] = None + geometry: list[Any] = field(default_factory=list) + created_at: Optional[str] = None + approved_at: Optional[str] = None + approved_by: Optional[str] = None + associated_file: Optional[str] = None + file: Optional[str] = None + deleted: bool = False + deleted_at: Optional[str] = None + deleted_by: Optional[str] = None + created_by_model: Optional[str] = None + old_geometry: Optional[Any] = None + set_name: Optional[str] = None + resource_filename: Optional[str] = None + resource_modality: Optional[str] = None + annotation_worklist_name: Optional[str] = None + user_info: Optional[dict[str, str]] = None + values: Optional[Any] = None + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> Annotation: + """ + Create an Annotation instance from a dictionary. + + Args: + data: Dictionary containing annotation data from API + + Returns: + Annotation instance + """ + + + # Convert field names and filter valid fields + converted_data = {} + for key, value in data.items(): + # Map field names if needed + mapped_key = _FIELD_MAPPING.get(key, key) + converted_data[mapped_key] = value + + if 'scope' not in converted_data: + converted_data['scope'] = 'image' if converted_data.get('frame_index') is None else 'frame' + + if converted_data['annotation_type'] in ['segmentation']: + if converted_data.get('file') is None: + raise ValueError(f"Segmentation annotations must have an associated file. {data}") + + # Create instance with only valid fields + valid_fields = {f.name for f in cls.__dataclass_fields__.values()} + filtered_data = {k: v for k, v in converted_data.items() if k in valid_fields} + + return cls(**filtered_data) + + def to_dict(self) -> dict[str, Any]: + """ + Convert the annotation to a dictionary format. + + Returns: + Dictionary representation of the annotation + """ + result = {} + for key, value in self.__dict__.items(): + # Handle special serialization cases + if isinstance(value, (np.ndarray, np.generic)): + value = value.tolist() + elif isinstance(value, datetime): + value = value.isoformat() + elif isinstance(value, Path): + value = str(value) + + result[key] = value + if 'file' not in result: + raise ValueError(f"Segmentation annotations must have an associated file. {self}") + return result + + @property + def name(self) -> str: + """Get the annotation name (alias for identifier).""" + return self.identifier + + @property + def type(self) -> str: + """Get the annotation type.""" + return self.annotation_type + + @property + def value(self) -> Optional[str]: + """Get the annotation value (for category annotations).""" + return self.text_value + + @property + def index(self) -> Optional[int]: + """Get the frame index (alias for frame_index).""" + return self.frame_index + + @property + def added_by(self) -> str: + """Get the creator email (alias for created_by).""" + return self.created_by + + # @property + # def file(self) -> Optional[str]: + # """Get the associated file path.""" + # return self.associated_file + + # @file.setter + # def file(self, value: Optional[str]) -> None: + # """Set the associated file path.""" + # self.associated_file = value + + def is_segmentation(self) -> bool: + """Check if this is a segmentation annotation.""" + return self.annotation_type == 'segmentation' + + def is_label(self) -> bool: + """Check if this is a label annotation.""" + return self.annotation_type == 'label' + + def is_category(self) -> bool: + """Check if this is a category annotation.""" + return self.annotation_type == 'category' + + def is_frame_scoped(self) -> bool: + """Check if this annotation is frame-scoped.""" + return self.scope == 'frame' + + def is_image_scoped(self) -> bool: + """Check if this annotation is image-scoped.""" + return self.scope == 'image' + + def get_created_datetime(self) -> Optional[datetime]: + """ + Get the creation datetime as a datetime object. + + Returns: + datetime object or None if created_at is not set + """ + if self.created_at: + try: + return datetime.fromisoformat(self.created_at.replace('Z', '+00:00')) + except ValueError: + _LOGGER.warning(f"Could not parse created_at datetime: {self.created_at}") + return None + + def __repr__(self) -> str: + """String representation of the annotation.""" + return (f"Annotation(id='{self.id}', identifier='{self.identifier}', " + f"type='{self.annotation_type}', scope='{self.scope}', resource_id='{self.resource_id}')") diff --git a/datamint/dataset/base_dataset.py b/datamint/dataset/base_dataset.py index 75a3369b..de0053fe 100644 --- a/datamint/dataset/base_dataset.py +++ b/datamint/dataset/base_dataset.py @@ -15,9 +15,12 @@ from torch import Tensor from datamint.apihandler.base_api_handler import DatamintException from medimgkit.dicom_utils import is_dicom -import cv2 from medimgkit.io_utils import read_array_normalized from datetime import datetime +from pathlib import Path +from mimetypes import guess_extension +from datamint.dataset.annotation import Annotation +import cv2 _LOGGER = logging.getLogger(__name__) @@ -80,10 +83,10 @@ def __init__( exclude_frame_label_names: Optional[list[str]] = None, ): self._validate_inputs(project_name, include_annotators, exclude_annotators, - include_segmentation_names, exclude_segmentation_names, - include_image_label_names, exclude_image_label_names, - include_frame_label_names, exclude_frame_label_names) - + include_segmentation_names, exclude_segmentation_names, + include_image_label_names, exclude_image_label_names, + include_frame_label_names, exclude_frame_label_names) + self._initialize_config( project_name, auto_update, all_annotations, return_dicom, return_metainfo, return_annotations, return_frame_by_frame, @@ -92,7 +95,7 @@ def __init__( include_image_label_names, exclude_image_label_names, include_frame_label_names, exclude_frame_label_names ) - + self._setup_api_handler(server_url, api_key, auto_update) self._setup_directories(root) self._setup_dataset() @@ -121,7 +124,7 @@ def _validate_inputs( (include_image_label_names, exclude_image_label_names, "image_label_names"), (include_frame_label_names, exclude_frame_label_names, "frame_label_names"), ] - + for include_param, exclude_param, param_name in filter_pairs: if include_param is not None and exclude_param is not None: raise ValueError(f"Cannot set both include_{param_name} and exclude_{param_name} at the same time") @@ -167,13 +170,14 @@ def _initialize_config( # Internal state self.__logged_uint16_conversion = False + self.auto_update = auto_update 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, + root_url=server_url, api_key=api_key, check_connection=auto_update ) @@ -206,42 +210,68 @@ def _setup_directories(self, root: str | None) -> None: self.dataset_dir = os.path.join(root, self.project_name) self.dataset_zippath = os.path.join(root, f'{self.project_name}.zip') + if not os.path.exists(self.dataset_dir): + os.makedirs(self.dataset_dir, exist_ok=True) + os.makedirs(os.path.join(self.dataset_dir, 'images'), exist_ok=True) + os.makedirs(os.path.join(self.dataset_dir, 'masks'), exist_ok=True) + def _setup_dataset(self) -> None: """Setup dataset by downloading or loading existing data.""" - local_dataset_exists = os.path.exists(os.path.join(self.dataset_dir, 'dataset.json')) - - if local_dataset_exists and not hasattr(self, 'project_info'): - # We might not need project info if not updating - self.dataset_id = None - else: - self.project_info = self.get_info() - self.dataset_id = self.project_info['dataset_id'] - - self._handle_dataset_download_or_update(local_dataset_exists) - self._load_metadata() + self._server_dataset_info = None + local_load_success = self._load_metadata() + self._handle_dataset_download_or_update(local_load_success) + self._apply_annotation_filters() - def _handle_dataset_download_or_update(self, local_dataset_exists: bool) -> None: + def _handle_dataset_download_or_update(self, local_load_success: bool) -> None: """Handle dataset download or update logic.""" - if local_dataset_exists: - _LOGGER.info(f"Dataset directory already exists: {self.dataset_dir}") + + if local_load_success: + _LOGGER.debug(f"Dataset directory already exists: {self.dataset_dir}") # Check for updates if auto_update is enabled and we have API access - if hasattr(self, 'project_info'): + if self.auto_update: _LOGGER.info("Checking for updates...") self._check_version() else: - if self.api_key is None: - raise DatamintDatasetException("API key is required to download the dataset.") - _LOGGER.info(f"No data found at {self.dataset_dir}. Downloading...") - self.download_project() - - def _load_metadata(self) -> None: + self._check_version() + # if self.api_key is None: + # raise DatamintDatasetException("API key is required to download the dataset.") + # self.project_info = self.get_info() + # self.dataset_id = self.project_info['dataset_id'] + # _LOGGER.info(f"No data found at {self.dataset_dir}. Downloading...") + # self.download_project() + + def _load_metadata(self) -> bool: """Load and process dataset metadata.""" - if not hasattr(self, 'metainfo'): - with open(os.path.join(self.dataset_dir, 'dataset.json'), 'r') as file: + if hasattr(self, 'metainfo'): + _LOGGER.warning("Metadata already loaded.") + metadata_path = os.path.join(self.dataset_dir, 'dataset.json') + 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['updated_at'] = None + self.metainfo['resources'] = [] + self.metainfo['all_annotations'] = self.all_annotations + self.images_metainfo = self.metainfo['resources'] + return False + else: + with open(metadata_path, 'r') as file: self.metainfo = json.load(file) - self.images_metainfo = self.metainfo['resources'] - self._apply_annotation_filters() + # Convert annotations from dict to Annotation objects + self._convert_metainfo_to_clsobj() + return True + + def _convert_metainfo_to_clsobj(self): + for imginfo in self.images_metainfo: + if 'annotations' in imginfo: + for ann in imginfo['annotations']: + if 'resource_id' not in ann: + ann['resource_id'] = imginfo['id'] + if 'id' not in ann: + ann['id'] = None + imginfo['annotations'] = [Annotation.from_dict(ann) if isinstance(ann, dict) else ann + for ann in imginfo['annotations']] def _apply_annotation_filters(self) -> None: """Apply annotation filters and remove unannotated images if needed.""" @@ -261,7 +291,7 @@ def _post_process_data(self) -> None: self._calculate_dataset_length() self._precompute_frame_data() self._setup_labels() - + if self.discard_without_annotations and self.return_frame_by_frame: self._filter_unannotated() @@ -298,13 +328,13 @@ def _filter_unannotated(self) -> None: annotations = item_meta.get('annotations', []) # Check if there are any segmentation annotations - has_segmentations = any(ann['type'] == 'segmentation' for ann in annotations) + has_segmentations = any(ann.type == 'segmentation' for ann in annotations) if has_segmentations: filtered_indices.append(self.subset_indices[i]) self.subset_indices = filtered_indices - _LOGGER.info(f"Filtered dataset: {len(self.subset_indices)} frames with segmentations") + _LOGGER.debug(f"Filtered dataset: {len(self.subset_indices)} frames with segmentations") def __compute_num_frames_per_resource(self) -> list[int]: """Compute number of frames for each resource.""" @@ -340,10 +370,10 @@ def segmentation_labels_set(self) -> list[str]: def _get_annotations_internal( self, - annotations: list[dict], + annotations: list[Annotation], type: Literal['label', 'category', 'segmentation', 'all'] = 'all', scope: Literal['frame', 'image', 'all'] = 'all' - ) -> list[dict]: + ) -> list[Annotation]: """Internal method to filter annotations by type and scope.""" if type not in ['label', 'category', 'segmentation', 'all']: raise ValueError(f"Invalid value for 'type': {type}") @@ -352,14 +382,14 @@ def _get_annotations_internal( filtered_annotations = [] for ann in annotations: - ann_scope = 'image' if ann.get('index', None) is None else 'frame' - - type_matches = type == 'all' or ann['type'] == type + ann_scope = 'image' if ann.index is None else 'frame' + + type_matches = type == 'all' or ann.type == type scope_matches = scope == 'all' or scope == ann_scope - + if type_matches and scope_matches: filtered_annotations.append(ann) - + return filtered_annotations def get_annotations( @@ -367,7 +397,7 @@ def get_annotations( index: int, type: Literal['label', 'category', 'segmentation', 'all'] = 'all', scope: Literal['frame', 'image', 'all'] = 'all' - ) -> list[dict]: + ) -> list[Annotation]: """Returns the annotations of the image at the given index. Args: @@ -380,7 +410,7 @@ def get_annotations( """ if index >= len(self): raise IndexError(f"Index {index} out of bounds for dataset of length {len(self)}") - + imginfo = self._get_image_metainfo(index) return self._get_annotations_internal(imginfo['annotations'], type=type, scope=scope) @@ -404,7 +434,7 @@ def read_number_of_frames(filepath: str) -> int: def get_resources_ids(self) -> list[str]: """Get list of resource IDs.""" return [ - self.__getitem_internal(i, only_load_metainfo=True)['metainfo']['id'] + self.__getitem_internal(i, only_load_metainfo=True)['metainfo']['id'] for i in self.subset_indices ] @@ -426,13 +456,13 @@ def _get_labels_set(self, framed: bool) -> tuple[dict, dict[str, dict[str, int]] for i in range(len(self)): # Collect labels by type label_anns = self.get_annotations(i, type='label', scope=scope) - multilabel_set.update(ann['name'] for ann in label_anns) + multilabel_set.update(ann.name for ann in label_anns) seg_anns = self.get_annotations(i, type='segmentation', scope=scope) - segmentation_labels.update(ann['name'] for ann in seg_anns) + segmentation_labels.update(ann.name for ann in seg_anns) cat_anns = self.get_annotations(i, type='category', scope=scope) - multiclass_set.update((ann['name'], ann['value']) for ann in cat_anns) + multiclass_set.update((ann.name, ann.value) for ann in cat_anns) # Sort and create mappings multilabel_list = sorted(multilabel_set) @@ -444,13 +474,13 @@ def _get_labels_set(self, framed: bool) -> tuple[dict, dict[str, dict[str, int]] 'segmentation': segmentation_list, 'multiclass': multiclass_list } - + codes_map = { 'multilabel': {label: idx for idx, label in enumerate(multilabel_list)}, 'segmentation': {label: idx + 1 for idx, label in enumerate(segmentation_list)}, 'multiclass': {label: idx for idx, label in enumerate(multiclass_list)} } - + return sets, codes_map def get_framelabel_distribution(self, normalize: bool = False) -> dict[str, float]: @@ -471,23 +501,23 @@ def _get_label_distribution(self, ann_type: str, scope: str, normalize: bool) -> raise ValueError(f"Unsupported combination: type={ann_type}, scope={scope}") distribution = {label: 0 for label in labels} - + for imginfo in self.images_metainfo: for ann in imginfo.get('annotations', []): condition_met = ( - ann['type'] == ann_type and - (scope == 'all' or - (scope == 'frame' and ann.get('index') is not None) or - (scope == 'image' and ann.get('index') is None)) + ann.type == ann_type and + (scope == 'all' or + (scope == 'frame' and ann.index is not None) or + (scope == 'image' and ann.index is None)) ) - if condition_met and ann['name'] in distribution: - distribution[ann['name']] += 1 + if condition_met and ann.name in distribution: + distribution[ann.name] += 1 if normalize: total = sum(distribution.values()) if total > 0: distribution = {k: v / total for k, v in distribution.items()} - + return distribution def _check_integrity(self) -> None: @@ -497,16 +527,19 @@ def _check_integrity(self) -> None: filepath = os.path.join(self.dataset_dir, imginfo['file']) if not os.path.isfile(filepath): missing_files.append(imginfo['file']) - + if missing_files: raise DatamintDatasetException(f"Image files not found: {missing_files}") def _get_datasetinfo(self) -> dict: """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() for dataset in all_datasets: if dataset['id'] == self.dataset_id: + self._server_dataset_info = dataset return dataset available_datasets = [(d['name'], d['id']) for d in all_datasets] @@ -517,6 +550,8 @@ def _get_datasetinfo(self) -> dict: def get_info(self) -> dict: """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) if 'error' in project: available_projects = project['all_projects'] @@ -524,6 +559,8 @@ def get_info(self) -> dict: f"Project with name '{self.project_name}' not found. " f"Available projects: {available_projects}" ) + self.project_info = project + self.dataset_id = project['dataset_id'] return project def _run_request(self, session, request_args) -> requests.Response: @@ -533,62 +570,11 @@ def _run_request(self, session, request_args) -> requests.Response: response.raise_for_status() return response - def _get_jwttoken(self, dataset_id, session) -> str: - if dataset_id is None: - raise ValueError("Dataset ID is required to download the dataset.") - request_params = { - 'method': 'GET', - 'url': f'{self.server_url}/datasets/{dataset_id}/download/png', - 'headers': {'apikey': self.api_key}, - 'stream': True - } - _LOGGER.debug(f"Getting jwt token for dataset {dataset_id}...") - response = self._run_request(session, request_params) - progress_bar = None - number_processed_images = 0 - - # check if the response is a stream of data and everything is ok - if response.status_code != 200: - msg = f"Getting jwt token failed with status code={response.status_code}: {response.text}" - raise DatamintDatasetException(msg) - - try: - response_iterator = response.iter_lines(decode_unicode=True) - for line in response_iterator: - line = line.strip() - if 'event: error' in line: - error_msg = line+'\n' - error_msg += '\n'.join(response_iterator) - raise DatamintDatasetException(f"Getting jwt token failed:\n{error_msg}") - if not line.startswith('data:'): - continue - dataline = yaml.safe_load(line)['data'] - if 'zip' in dataline: - _LOGGER.debug(f"Got jwt token for dataset {dataset_id}") - return dataline['zip'] # Function normally ends here - elif 'processedImages' in dataline: - if progress_bar is None: - total_size = int(dataline['totalImages']) - progress_bar = tqdm(total=total_size, unit='iB', unit_scale=True) - processed_images = int(dataline['processedImages']) - if number_processed_images < processed_images: - progress_bar.update(processed_images - number_processed_images) - number_processed_images = processed_images - else: - _LOGGER.warning(f"Unknown data line: {dataline}") - except Exception as e: - raise e - finally: - if progress_bar is not None: - progress_bar.close() - - raise DatamintDatasetException("Getting jwt token failed! No dataline with 'zip' entry found.") - def __repr__(self) -> str: """String representation of the dataset.""" head = f"Dataset {self.project_name}" body = [f"Number of datapoints: {self.__len__()}"] - + if self.root is not None: body.append(f"Location: {self.dataset_dir}") @@ -603,7 +589,7 @@ def __repr__(self) -> str: (self.include_frame_label_names, "Including only frame labels"), (self.exclude_frame_label_names, "Excluding frame labels"), ] - + for filter_value, description in filter_info: if filter_value is not None: body.append(f"{description}: {filter_value}") @@ -613,7 +599,6 @@ def __repr__(self) -> str: def download_project(self) -> None: """Download project data from API.""" - from torchvision.datasets.utils import extract_archive dataset_info = self._get_datasetinfo() self.dataset_id = dataset_info['id'] @@ -625,34 +610,43 @@ def download_project(self) -> None: 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'] + return self.dataset_id + def _extract_and_update_metadata(self) -> None: """Extract downloaded archive and update metadata.""" from torchvision.datasets.utils import extract_archive - + if os.path.exists(self.dataset_dir): _LOGGER.info(f"Deleting existing dataset directory: {self.dataset_dir}") shutil.rmtree(self.dataset_dir) - + extract_archive(self.dataset_zippath, self.dataset_dir, remove_finished=True) - + # Load and update metadata datasetjson_path = os.path.join(self.dataset_dir, 'dataset.json') with open(datasetjson_path, 'r') as file: self.metainfo = json.load(file) - + self._update_metadata_timestamps() - + # Save updated metadata with open(datasetjson_path, 'w') as file: - json.dump(self.metainfo, file) + json.dump(self.metainfo, file, default=lambda o: o.to_dict() if hasattr(o, 'to_dict') else o) + + self.images_metainfo = self.metainfo['resources'] + # self._convert_metainfo_to_clsobj() def _update_metadata_timestamps(self) -> None: """Update metadata with correct timestamps.""" @@ -662,7 +656,7 @@ def _update_metadata_timestamps(self) -> None: try: local_time = datetime.fromisoformat(self.metainfo['updated_at']) server_time = datetime.fromisoformat(self.last_updaded_at) - + if local_time < server_time: _LOGGER.warning( f"Inconsistent updated_at dates detected " @@ -702,7 +696,7 @@ def _process_image_array(self, img: np.ndarray) -> Tensor: img = img.astype(np.uint8) img_tensor = torch.from_numpy(img).contiguous() - + if isinstance(img_tensor, torch.ByteTensor): img_tensor = img_tensor.to(dtype=torch.get_default_dtype()).div(255) @@ -712,18 +706,18 @@ def _get_image_metainfo(self, index: int, bypass_subset_indices: bool = False) - """Get metadata for image at given index.""" if not bypass_subset_indices: index = self.subset_indices[index] - + if self.return_frame_by_frame: resource_id, frame_index = self.__find_index(index) img_metainfo = dict(self.images_metainfo[resource_id]) # Copy img_metainfo['frame_index'] = frame_index img_metainfo['annotations'] = [ ann for ann in img_metainfo['annotations'] - if ann['index'] is None or ann['index'] == frame_index + if ann.index is None or ann.index == frame_index ] else: img_metainfo = self.images_metainfo[index] - + return img_metainfo def __find_index(self, index: int) -> tuple[int, int]: @@ -733,7 +727,7 @@ def __find_index(self, index: int) -> tuple[int, int]: return resource_index, frame_index def __getitem_internal( - self, + self, index: int, only_load_metainfo: bool = False ) -> dict[str, Tensor | FileDataset | dict | list]: @@ -743,7 +737,7 @@ def __getitem_internal( else: resource_index = index frame_idx = None - + img_metainfo = self._get_image_metainfo(index, bypass_subset_indices=True) if only_load_metainfo: @@ -755,9 +749,9 @@ def __getitem_internal( return self._build_item_dict(img, ds, img_metainfo) def _build_item_dict( - self, - img: Tensor, - ds: FileDataset | None, + self, + img: Tensor, + ds: FileDataset | None, img_metainfo: dict ) -> dict[str, Any]: """Build the return dictionary for __getitem__.""" @@ -772,7 +766,7 @@ def _build_item_dict( return ret - def _filter_annotations(self, annotations: list[dict]) -> list[dict]: + def _filter_annotations(self, annotations: list[Annotation]) -> list[Annotation]: """Filter annotations based on the filtering settings.""" if annotations is None: return [] @@ -785,19 +779,19 @@ def _filter_annotations(self, annotations: list[dict]) -> list[dict]: return filtered_annotations - def _should_include_annotation(self, ann: dict) -> bool: + def _should_include_annotation(self, ann: Annotation) -> bool: """Check if an annotation should be included based on all filters.""" - if not self._should_include_annotator(ann['added_by']): + if not self._should_include_annotator(ann.created_by): return False - if ann['type'] == 'segmentation': - return self._should_include_segmentation(ann['name']) - elif ann['type'] == 'label': - if ann.get('index', None) is None: - return self._should_include_image_label(ann['name']) + if ann.type == 'segmentation': + return self._should_include_segmentation(ann.name) + elif ann.type == 'label': + if ann.index is None: + return self._should_include_image_label(ann.name) else: - return self._should_include_frame_label(ann['name']) - + return self._should_include_frame_label(ann.name) + return True def __getitem__(self, index: int) -> dict[str, Tensor | FileDataset | dict | list]: @@ -825,16 +819,17 @@ def __len__(self) -> int: def _check_version(self) -> None: """Check if local dataset version is up to date.""" - metainfo_path = os.path.join(self.dataset_dir, 'dataset.json') - if not os.path.exists(metainfo_path): - self.download_project() - return - - with open(metainfo_path, 'r') as file: - local_dataset_info = json.load(file) - - local_updated_at = local_dataset_info.get('updated_at', None) - local_all_annotations = local_dataset_info.get('all_annotations', None) + # metainfo_path = os.path.join(self.dataset_dir, 'dataset.json') + # if not os.path.exists(metainfo_path): + # self.download_project() + # return + + if not hasattr(self, 'project_info'): + self.project_info = self.get_info() + self.dataset_id = self.project_info['dataset_id'] + + local_updated_at = self.metainfo.get('updated_at', None) + local_all_annotations = self.metainfo.get('all_annotations', None) try: external_metadata_info = self._get_datasetinfo() @@ -848,21 +843,198 @@ def _check_version(self) -> None: annotations_changed = local_all_annotations != self.all_annotations version_outdated = local_updated_at is None or local_updated_at < server_updated_at - if annotations_changed or version_outdated: - if annotations_changed: - _LOGGER.info( - f"The 'all_annotations' parameter has changed. " - f"Previous: {local_all_annotations}, Current: {self.all_annotations}." - ) - else: - _LOGGER.info( - f"A newer version of the dataset is available. " - f"Your version: {local_updated_at}. Last version: {server_updated_at}." - ) - self.download_project() + if annotations_changed: + _LOGGER.info( + f"The 'all_annotations' parameter has changed. " + f"Previous: {local_all_annotations}, Current: {self.all_annotations}." + ) + # self.download_project() + self._incremental_update() + elif version_outdated: + _LOGGER.info( + f"A newer version of the dataset is available. " + f"Your version: {local_updated_at}. Last version: {server_updated_at}." + ) + self._incremental_update() else: _LOGGER.info('Local version is up to date with the latest version.') + def _fetch_new_resources(self, + all_uptodate_resources: list[dict]) -> 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: + 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]: + local_resources = self.images_metainfo + all_uptodate_resources_ids = [res['id'] for res in all_uptodate_resources] + deleted_resources = [] + for resource in local_resources: + try: + res_idx = all_uptodate_resources_ids.index(resource['id']) + if resource.get('deleted_at', None): # was deleted on server + if local_resources[res_idx].get('deleted_at_local', None) is None: + deleted_resources.append(resource) + except ValueError: + deleted_resources.append(resource) + + return deleted_resources + + def _incremental_update(self) -> None: + # local_updated_at = self.metainfo.get('updated_at', None) + # external_metadata_info = self._get_datasetinfo() + # server_updated_at = external_metadata_info['updated_at'] + + ### RESOURCES ### + all_uptodate_resources = self.api_handler.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) + + for r in new_resources: + self._new_resource_created(r) + 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...") + self.api_handler.download_multiple_resources(new_resources_ids, + save_path=new_resources_path) + _LOGGER.info(f"Downloaded {len(new_resources)} new resources.") + + for r in deleted_resources: + self._resource_deleted(r) + ################ + + ### ANNOTATIONS ### + all_annotations = self.api_handler.get_annotations(worklist_id=self.project_info['worklist_id'], + status='published' if self.all_annotations else None) + # group annotations by resource ID + annotations_by_resource = {} + 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'] + if resource_id not in annotations_by_resource: + annotations_by_resource[resource_id] = [] + annotations_by_resource[resource_id].append(ann) + + # update annotations in resources + for resource in tqdm(self.images_metainfo, desc="Updating annotations"): + resource_id = resource['id'] + new_resource_annotations = annotations_by_resource.get(resource_id, []) + old_resource_annotations = resource.get('annotations', []) + + # 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]) + + # 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] + annotations_to_remove = [ann for ann in old_resource_annotations + if getattr(ann, 'id', 'NA') not in new_ann_ids] + + for ann in annotations_to_add: + filepath = self._get_annotation_file_path(ann) + if filepath is not None: # None means it is not a segmentation + # download + filepath = Path(self.dataset_dir) / filepath + filepath.parent.mkdir(parents=True, exist_ok=True) + self.api_handler.download_segmentation_file(ann, fpath_out=filepath) + + # Process annotation changes + for ann in annotations_to_remove: + filepath = getattr(ann, 'file', None) if hasattr(ann, 'file') else ann.get('file', None) + if filepath is None: + # Not a segmentation annotation + continue + + try: + filepath = Path(self.dataset_dir) / filepath + # delete the local annotation file if it exists + if filepath.exists(): + os.remove(filepath) + except Exception as e: + _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] + ################### + # 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) + + def _get_resource_file_path(self, resource: dict) -> Path: + """Get the local file path for a resource.""" + if 'file' in resource and resource['file'] is not None: + return Path(resource['file']) + else: + ext = guess_extension(resource['mimetype'], strict=False) + if ext is None: + _LOGGER.warning(f"Could not guess extension for resource {resource['id']}.") + ext = '' + return Path('images', f"{resource['id']}{ext}") + + def _get_annotation_file_path(self, annotation: dict | Annotation) -> Path | None: + """Get the local file path for an annotation.""" + if isinstance(annotation, Annotation): + if annotation.file: + return Path(annotation.file) + elif annotation.type == 'segmentation': + return Path('masks', + annotation.created_by, + annotation.resource_id, + annotation.id) + else: + # Handle dict format for backwards compatibility + if 'file' in annotation: + return Path(annotation['file']) + elif annotation.get('annotation_type', annotation.get('type')) == 'segmentation': + return Path('masks', + annotation['created_by'], + annotation['resource_id'], + annotation['id']) + return None + + def _new_resource_created(self, resource: dict) -> None: + """Handle a new resource created in the dataset.""" + if 'annotations' not in resource: + resource['annotations'] = [] # Initialize as empty list for Annotation objects + self.images_metainfo.append(resource) + + if hasattr(self, 'num_frames_per_resource'): + raise NotImplementedError('Cannot handle new resources after dataset initialization') + + def _resource_deleted(self, resource: dict) -> None: + """Handle a resource deleted from the dataset.""" + + # remove from metadata + for i, imginfo in enumerate(self.images_metainfo): + if imginfo['id'] == resource['id']: + deleted_metainfo = self.images_metainfo.pop(i) + break + else: + _LOGGER.warning(f"Resource {resource['id']} not found in dataset metadata.") + return + + # delete from system file + if os.path.exists(deleted_metainfo['file']): + os.remove(deleted_metainfo['file']) + + # delete associated annotations + for ann in deleted_metainfo.get('annotations', []): + ann_file = getattr(ann, 'file', None) if hasattr(ann, 'file') else ann.get('file', None) + if ann_file is not None: + os.remove(ann_file) + def __add__(self, other): """Concatenate datasets.""" from torch.utils.data import ConcatDataset @@ -885,13 +1057,13 @@ def get_collate_fn(self) -> Callable: def collate_fn(batch: list[dict]) -> dict: if not batch: return {} - + keys = batch[0].keys() collated_batch = {} - + for key in keys: values = [item[key] for item in batch] - + if isinstance(values[0], torch.Tensor): shapes = [tensor.shape for tensor in values] if all(shape == shapes[0] for shape in shapes): diff --git a/datamint/dataset/dataset.py b/datamint/dataset/dataset.py index 8ff5293b..7d500e8c 100644 --- a/datamint/dataset/dataset.py +++ b/datamint/dataset/dataset.py @@ -7,6 +7,7 @@ import logging from PIL import Image import albumentations +from datamint.dataset.annotation import Annotation _LOGGER = logging.getLogger(__name__) @@ -116,12 +117,12 @@ def __init__(self, if semantic_seg_merge_strategy is not None and not return_as_semantic_segmentation: raise ValueError("semantic_seg_merge_strategy can only be used if return_as_semantic_segmentation is True") - def _load_segmentations(self, annotations: list[dict], img_shape) -> tuple[dict[str, list], dict[str, list]]: + def _load_segmentations(self, annotations: list[Annotation], img_shape) -> tuple[dict[str, list], dict[str, list]]: """ Load segmentations from annotations. Args: - annotations: list of annotations. Each annotation is a dictionary with keys 'type', 'file', 'added_by', 'name', 'index'. + annotations: list of Annotation objects img_shape: shape of the image (#frames, C, H, W) Returns: @@ -142,14 +143,14 @@ def _load_segmentations(self, annotations: list[dict], img_shape) -> tuple[dict[ # Load segmentation annotations for ann in annotations: - if ann['type'] != 'segmentation': + if ann.type != 'segmentation': continue - if 'file' not in ann: + if ann.file is None: _LOGGER.warning(f"Segmentation annotation without file in annotations {ann}") continue - author = ann['added_by'] + author = ann.added_by - segfilepath = ann['file'] # png file + segfilepath = ann.file # png file segfilepath = os.path.join(self.dataset_dir, segfilepath) # FIXME: avoid enforcing resizing the mask seg = (Image.open(segfilepath) @@ -161,11 +162,11 @@ def _load_segmentations(self, annotations: list[dict], img_shape) -> tuple[dict[ seg = torch.from_numpy(seg) seg = seg == 255 # binary mask # map the segmentation label to the code - seg_code = self.frame_lcodes['segmentation'][ann['name']] + seg_code = self.frame_lcodes['segmentation'][ann.name] if self.return_frame_by_frame: frame_index = 0 else: - frame_index = ann['index'] + frame_index = ann.index if author not in segmentations.keys(): segmentations[author] = [None] * nframes @@ -476,14 +477,14 @@ def __getitem__(self, index) -> dict[str, Any]: return new_item def _convert_labels_annotations(self, - annotations: list[dict], + annotations: list[Annotation], num_frames: int = None) -> dict[str, torch.Tensor]: """ Converts the annotations, of the same type and scope, to tensor of shape (num_frames, num_labels) for each annotator. Args: - annotations: list of annotations + annotations: list of Annotation objects num_frames: number of frames in the video Returns: @@ -505,14 +506,14 @@ def _convert_labels_annotations(self, if len(annotations) == 0: return frame_labels_byuser for ann in annotations: - user_id = ann['added_by'] + user_id = ann.added_by - frame_idx = ann.get('index', None) + frame_idx = ann.index if user_id not in frame_labels_byuser.keys(): frame_labels_byuser[user_id] = torch.zeros(size=labels_ret_size, dtype=torch.int32) labels_onehot_i = frame_labels_byuser[user_id] - code = label2code[ann['name']] + code = label2code[ann.name] if frame_idx is None: labels_onehot_i[code] = 1 else: From d4bcc8faa847e8a70488c40e39e83efaa6f06cb8 Mon Sep 17 00:00:00 2001 From: Lucashsmello Date: Fri, 1 Aug 2025 16:50:36 -0300 Subject: [PATCH 2/5] Implement asynchronous batch downloading of segmentation files in AnnotationAPIHandler and update DatamintBaseDataset to utilize this feature. Bump version to 1.7.0. --- datamint/apihandler/annotation_api_handler.py | 67 +++++++++++++++++++ datamint/dataset/base_dataset.py | 18 ++++- pyproject.toml | 2 +- 3 files changed, 83 insertions(+), 4 deletions(-) diff --git a/datamint/apihandler/annotation_api_handler.py b/datamint/apihandler/annotation_api_handler.py index ec182bb3..5cbf3a4e 100644 --- a/datamint/apihandler/annotation_api_handler.py +++ b/datamint/apihandler/annotation_api_handler.py @@ -15,6 +15,7 @@ import json from deprecated import deprecated from pathlib import Path +from tqdm.auto import tqdm _LOGGER = logging.getLogger(__name__) _USER_LOGGER = logging.getLogger('user_logger') @@ -1179,3 +1180,69 @@ def set_annotation_status(self, } resp = self._run_request(request_params) self._check_errors_response_json(resp) + + + async def _async_download_segmentation_file(self, + annotation: str | dict, + save_path: str | Path, + session: aiohttp.ClientSession | None = None, + progress_bar: tqdm | None = None): + """ + Asynchronously download a segmentation file. + + Args: + annotation (str | dict): The annotation unique id or an annotation object. + save_path (str | Path): The path to save the file. + session (aiohttp.ClientSession): The aiohttp session to use for the request. + progress_bar (tqdm | None): Optional progress bar to update after download completion. + """ + if isinstance(annotation, dict): + annotation_id = annotation['id'] + resource_id = annotation['resource_id'] + else: + annotation_id = annotation + # TODO: This is inefficient as it requires an extra API call per annotation + # Consider passing resource_id separately or caching annotation info + resource_id = self.get_annotation_by_id(annotation_id)['resource_id'] + + url = f'{self.root_url}/annotations/{resource_id}/annotations/{annotation_id}/file' + request_params = { + 'method': 'GET', + 'url': url + } + + try: + data_bytes = await self._run_request_async(request_params, session, 'content') + with open(save_path, 'wb') as f: + f.write(data_bytes) + if progress_bar: + progress_bar.update(1) + except ResourceNotFoundError as e: + e.set_params('annotation', {'annotation_id': annotation_id}) + raise e + + def download_multiple_segmentations(self, + annotations: list[str | dict], + save_paths: list[str | Path] | str + ) -> None: + """ + Download multiple segmentation files and save them to the specified paths. + + Args: + annotations (list[str | dict]): A list of annotation unique ids or annotation objects. + save_paths (list[str | Path] | str): A list of paths to save the files or a directory path. + """ + async def _download_all_async(): + async with aiohttp.ClientSession() as session: + tasks = [ + self._async_download_segmentation_file(annotation, save_path=path, session=session, progress_bar=progress_bar) + for annotation, path in zip(annotations, save_paths) + ] + await asyncio.gather(*tasks) + + if isinstance(save_paths, str): + save_paths = [os.path.join(save_paths, f"{ann['id'] if isinstance(ann, dict) else ann}") for ann in annotations] + + with tqdm(total=len(annotations), desc="Downloading segmentations", unit="file") as progress_bar: + loop = asyncio.get_event_loop() + loop.run_until_complete(_download_all_async()) diff --git a/datamint/dataset/base_dataset.py b/datamint/dataset/base_dataset.py index de0053fe..3c5137b4 100644 --- a/datamint/dataset/base_dataset.py +++ b/datamint/dataset/base_dataset.py @@ -924,8 +924,12 @@ def _incremental_update(self) -> None: annotations_by_resource[resource_id] = [] annotations_by_resource[resource_id].append(ann) + # Collect all segmentation annotations that need to be downloaded + segmentations_to_download = [] + segmentation_paths = [] + # update annotations in resources - for resource in tqdm(self.images_metainfo, desc="Updating annotations"): + for resource in self.images_metainfo: resource_id = resource['id'] new_resource_annotations = annotations_by_resource.get(resource_id, []) old_resource_annotations = resource.get('annotations', []) @@ -944,10 +948,11 @@ def _incremental_update(self) -> None: for ann in annotations_to_add: filepath = self._get_annotation_file_path(ann) if filepath is not None: # None means it is not a segmentation - # download + # Collect for batch download filepath = Path(self.dataset_dir) / filepath filepath.parent.mkdir(parents=True, exist_ok=True) - self.api_handler.download_segmentation_file(ann, fpath_out=filepath) + segmentations_to_download.append(ann) + segmentation_paths.append(filepath) # Process annotation changes for ann in annotations_to_remove: @@ -966,6 +971,13 @@ def _incremental_update(self) -> None: # Update resource annotations list - convert to Annotation objects resource['annotations'] = [Annotation.from_dict(ann) for ann in 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) + _LOGGER.info(f"Downloaded {len(segmentations_to_download)} segmentation files.") + ################### # save updated metadata datasetjson_path = os.path.join(self.dataset_dir, 'dataset.json') diff --git a/pyproject.toml b/pyproject.toml index 69023bcd..541cbbd0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "datamint" description = "A library for interacting with the Datamint API, designed for efficient data management, processing and Deep Learning workflows." -version = "1.6.3-post1" +version = "1.7.0" dynamic = ["dependencies"] requires-python = ">=3.10" readme = "README.md" From 61a120a38d38d0a1e64bc08110f0250f5075407a Mon Sep 17 00:00:00 2001 From: Lucashsmello Date: Fri, 1 Aug 2025 17:19:05 -0300 Subject: [PATCH 3/5] fixes error when iterating dataset --- datamint/dataset/base_dataset.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/datamint/dataset/base_dataset.py b/datamint/dataset/base_dataset.py index 3c5137b4..1aa90d55 100644 --- a/datamint/dataset/base_dataset.py +++ b/datamint/dataset/base_dataset.py @@ -811,7 +811,8 @@ def __getitem__(self, index: int) -> dict[str, Tensor | FileDataset | dict | lis def __iter__(self): """Iterate over dataset items.""" for index in self.subset_indices: - yield self.__getitem_internal(index) + yield self.__getitem__(index) + # do not use __getitem_internal__ here, so subclass only need to implement __getitem__ def __len__(self) -> int: """Return dataset length.""" @@ -979,6 +980,9 @@ def _incremental_update(self) -> None: _LOGGER.info(f"Downloaded {len(segmentations_to_download)} segmentation files.") ################### + # update metadata + 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: From ceeceb58cea91580cdcba3c1bc5ba0acf48383ad Mon Sep 17 00:00:00 2001 From: Lucashsmello Date: Fri, 1 Aug 2025 17:25:31 -0300 Subject: [PATCH 4/5] Fixed file removal paths for deleted resources and annotations. --- datamint/dataset/base_dataset.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/datamint/dataset/base_dataset.py b/datamint/dataset/base_dataset.py index 1aa90d55..042d9627 100644 --- a/datamint/dataset/base_dataset.py +++ b/datamint/dataset/base_dataset.py @@ -897,14 +897,15 @@ def _incremental_update(self) -> None: new_resources = self._fetch_new_resources(all_uptodate_resources) deleted_resources = self._fetch_deleted_resources(all_uptodate_resources) - for r in new_resources: - self._new_resource_created(r) - 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...") - self.api_handler.download_multiple_resources(new_resources_ids, - save_path=new_resources_path) - _LOGGER.info(f"Downloaded {len(new_resources)} new resources.") + if new_resources: + for r in new_resources: + self._new_resource_created(r) + 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...") + self.api_handler.download_multiple_resources(new_resources_ids, + save_path=new_resources_path) + _LOGGER.info(f"Downloaded {len(new_resources)} new resources.") for r in deleted_resources: self._resource_deleted(r) @@ -1043,13 +1044,13 @@ def _resource_deleted(self, resource: dict) -> None: # delete from system file if os.path.exists(deleted_metainfo['file']): - os.remove(deleted_metainfo['file']) + os.remove(os.path.join(self.dataset_dir, deleted_metainfo['file'])) # delete associated annotations for ann in deleted_metainfo.get('annotations', []): ann_file = getattr(ann, 'file', None) if hasattr(ann, 'file') else ann.get('file', None) if ann_file is not None: - os.remove(ann_file) + os.remove(os.path.join(self.dataset_dir, ann_file)) def __add__(self, other): """Concatenate datasets.""" From 0dca0adb6f372983f776939deca0841c263a7e5d Mon Sep 17 00:00:00 2001 From: Lucashsmello Date: Fri, 1 Aug 2025 18:48:12 -0300 Subject: [PATCH 5/5] Fixed checking of annotation file --- datamint/dataset/annotation.py | 2 +- datamint/dataset/base_dataset.py | 6 ------ datamint/dataset/dataset.py | 4 ++-- 3 files changed, 3 insertions(+), 9 deletions(-) diff --git a/datamint/dataset/annotation.py b/datamint/dataset/annotation.py index 11c91537..87fbde7b 100644 --- a/datamint/dataset/annotation.py +++ b/datamint/dataset/annotation.py @@ -142,7 +142,7 @@ def to_dict(self) -> dict[str, Any]: value = str(value) result[key] = value - if 'file' not in result: + if self.annotation_type == 'segmentation' and 'file' not in result: raise ValueError(f"Segmentation annotations must have an associated file. {self}") return result diff --git a/datamint/dataset/base_dataset.py b/datamint/dataset/base_dataset.py index 042d9627..407e47e7 100644 --- a/datamint/dataset/base_dataset.py +++ b/datamint/dataset/base_dataset.py @@ -233,12 +233,6 @@ def _handle_dataset_download_or_update(self, local_load_success: bool) -> None: self._check_version() else: self._check_version() - # if self.api_key is None: - # raise DatamintDatasetException("API key is required to download the dataset.") - # self.project_info = self.get_info() - # self.dataset_id = self.project_info['dataset_id'] - # _LOGGER.info(f"No data found at {self.dataset_dir}. Downloading...") - # self.download_project() def _load_metadata(self) -> bool: """Load and process dataset metadata.""" diff --git a/datamint/dataset/dataset.py b/datamint/dataset/dataset.py index 7d500e8c..56198906 100644 --- a/datamint/dataset/dataset.py +++ b/datamint/dataset/dataset.py @@ -148,7 +148,7 @@ def _load_segmentations(self, annotations: list[Annotation], img_shape) -> tuple if ann.file is None: _LOGGER.warning(f"Segmentation annotation without file in annotations {ann}") continue - author = ann.added_by + author = ann.created_by segfilepath = ann.file # png file segfilepath = os.path.join(self.dataset_dir, segfilepath) @@ -506,7 +506,7 @@ def _convert_labels_annotations(self, if len(annotations) == 0: return frame_labels_byuser for ann in annotations: - user_id = ann.added_by + user_id = ann.created_by frame_idx = ann.index