def _upload_segment( self, segment: Segment, *, jobs: int = 1, skip_uploaded_files: bool = False, pbar: Tqdm, ) -> SegmentClient: segment_client = self.get_or_create_segment(segment.name) all_data: Iterator[Union[AuthData, Data]] = filter( lambda data: pbar.update_for_skip(not isinstance(data, RemoteData) ), segment, # type: ignore[arg-type] ) if not skip_uploaded_files: segment_filter = all_data else: done_set = set(segment_client.list_data_paths()) segment_filter = filter( lambda data: pbar.update_for_skip(data.target_remote_path not in done_set), all_data, ) multithread_upload( # pylint: disable=protected-access segment_client._upload_or_import_data, segment_filter, callback=segment_client._synchronize_upload_info, jobs=jobs, pbar=pbar, ) return segment_client
def _extract_all_data( source_frames: Iterator[Tuple[Frame, ULID]], pbar: Tqdm, ) -> FrameDataGenerator: for frame, frame_id in source_frames: for sensor_name, data in frame.items(): if isinstance(data, RemoteData): pbar.update() continue yield data, sensor_name, frame_id.str
def _extract_unuploaded_data( source_frames: Iterator[Tuple[Frame, ULID]], pbar: Tqdm, *, done_frames: Dict[float, Frame]) -> FrameDataGenerator: for frame, frame_id in source_frames: done_frame = done_frames.get(frame_id.timestamp().timestamp) if done_frame: frame_id = done_frame.frame_id for sensor_name, data in frame.items(): if isinstance(data, RemoteData): pbar.update() continue if (done_frame and sensor_name in done_frame and done_frame[sensor_name].path == data.target_remote_path): pbar.update() continue yield data, sensor_name, frame_id.str
def upload_segment( self, segment: Segment, *, jobs: int = 1, skip_uploaded_files: bool = False, quiet: bool = False, _is_cli: bool = False, ) -> SegmentClient: """Upload a :class:`~tensorbay.dataset.segment.Segment` to the dataset. This function will upload all info contains in the input :class:`~tensorbay.dataset.segment.Segment`, which includes: - Create a segment using the name of input Segment. - Upload all Data in the Segment to the dataset. Arguments: segment: The :class:`~tensorbay.dataset.segment.Segment` contains the information needs to be upload. jobs: The number of the max workers in multi-thread uploading method. skip_uploaded_files: True for skipping the uploaded files. quiet: Set to True to stop showing the upload process bar. _is_cli: Whether the method is called by CLI. Raises: Exception: When the upload got interrupted by Exception. Returns: The :class:`~tensorbay.client.segment.SegmentClient` used for uploading the data in the segment. """ self._status.check_authority_for_draft() try: with Tqdm(len(segment), disable=quiet) as pbar: return self._upload_segment( segment, jobs=jobs, skip_uploaded_files=skip_uploaded_files, pbar=pbar, ) except Exception: if _is_cli: logger.error( UPLOAD_SEGMENT_RESUME_TEMPLATE_CLI, self._status.draft_number, ) else: logger.error( UPLOAD_SEGMENT_RESUME_TEMPLATE_SDK, self._status.draft_number, self._status.draft_number, ) raise
def upload_segment( self, segment: FusionSegment, *, jobs: int = 1, skip_uploaded_files: bool = False, quiet: bool = False, ) -> FusionSegmentClient: """Upload a fusion segment object to the draft. This function will upload all info contains in the input :class:`~tensorbay.dataset.segment.FusionSegment`, which includes: - Create a segment using the name of input fusion segment object. - Upload all sensors in the segment to the dataset. - Upload all frames in the segment to the dataset. Arguments: segment: The :class:`~tensorbay.dataset.segment.FusionSegment`. jobs: The number of the max workers in multi-thread upload. skip_uploaded_files: Set it to True to skip the uploaded files. quiet: Set to True to stop showing the upload process bar. Raises: Exception: When the upload got interrupted by Exception. Returns: The :class:`~tensorbay.client.segment.FusionSegmentClient` used for uploading the data in the segment. """ self._status.check_authority_for_draft() try: with Tqdm(sum(len(frame) for frame in segment), disable=quiet) as pbar: return self._upload_segment( segment, jobs=jobs, skip_uploaded_files=skip_uploaded_files, pbar=pbar) except Exception: logger.error( UPLOAD_SEGMENT_RESUME_TEMPLATE_SDK, self._status.draft_number, self._status.draft_number, ) raise
def upload_dataset( self, dataset: Union[Dataset, FusionDataset], draft_number: Optional[int] = None, *, branch_name: Optional[str] = None, jobs: int = 1, skip_uploaded_files: bool = False, quiet: bool = False, ) -> DatasetClientType: """Upload a local dataset to TensorBay. This function will upload all information contains in the :class:`~tensorbay.dataset.dataset.Dataset` or :class:`~tensorbay.dataset.dataset.FusionDataset`, which includes: - Create a TensorBay dataset with the name and type of input local dataset. - Upload all :class:`~tensorbay.dataset.segment.Segment` or :class:`~tensorbay.dataset.segment.FusionSegment` in the dataset to TensorBay. Arguments: dataset: The :class:`~tensorbay.dataset.dataset.Dataset` or :class:`~tensorbay.dataset.dataset. FusionDataset` needs to be uploaded. draft_number: The draft number. branch_name: The branch name. jobs: The number of the max workers in multi-thread upload. skip_uploaded_files: Set it to True to skip the uploaded files. quiet: Set to True to stop showing the upload process bar. Returns: The :class:`~tensorbay.client.dataset.DatasetClient` or :class:`~tensorbay.client.dataset.FusionDatasetClient` bound with the uploaded dataset. Raises: ValueError: When uploading the dataset based on both draft number and branch name is not allowed. Exception: When Exception was raised during uploading dataset. """ dataset_client = self.get_dataset(dataset.name, isinstance(dataset, FusionDataset)) if draft_number and branch_name: raise ValueError( "Uploading the dataset based on both draft number and branch name is not allowed" ) if draft_number: dataset_client.checkout(draft_number=draft_number) else: target_branch_name = branch_name if branch_name else dataset_client.status.branch_name drafts = dataset_client.list_drafts(branch_name=target_branch_name) if drafts: dataset_client.checkout(draft_number=drafts[0].number) else: dataset_client.create_draft( 'Draft autogenerated by "GAS.upload_dataset"', branch_name=target_branch_name) try: if dataset.catalog: dataset_client.upload_catalog(dataset.catalog) dataset_client.update_notes( **dataset.notes) # type: ignore[arg-type] if isinstance(dataset, Dataset): data_count = sum(len(segment) for segment in dataset) else: data_count = sum( sum(len(frame) for frame in fusion_segment) for fusion_segment in dataset) with Tqdm(data_count, disable=quiet) as pbar: for segment in dataset: dataset_client._upload_segment( # pylint: disable=protected-access segment, # type: ignore[arg-type] jobs=jobs, skip_uploaded_files=skip_uploaded_files, pbar=pbar, ) except Exception: if draft_number: logger.error( UPLOAD_DATASET_RESUME_TEMPLATE, dataset_client.status.draft_number, dataset_client.status.draft_number, ) raise return dataset_client