def test_upload_segment(self, mocker):
self.fusion_dataset_client._status.checkout(draft_number=1)
segment_test = FusionSegment(name="test1")
for i in range(5):
temp_frame = Frame()
temp_ulid = from_timestamp(10 * i + 10)
temp_frame.frame_id = temp_ulid
if i % 2 == 0:
temp_frame["camera"] = Data(f"{i}.png")
else:
temp_frame["lidar"] = Data(f"{i}.png")
segment_test.append(temp_frame)
segment_client = FusionSegmentClient(
name="test1", data_client=self.fusion_dataset_client)
upload_segment = mocker.patch(
f"{dataset.__name__}.FusionDatasetClient._upload_segment",
return_value=segment_client)
assert self.fusion_dataset_client.upload_segment(
segment_test).name == "test1"
args, keywords = upload_segment.call_args
assert args[0] == segment_test
assert keywords["jobs"] == 1
assert not keywords["skip_uploaded_files"]
upload_segment.assert_called_once()
def _generate_segments(
self,
offset: int = 0,
limit: int = 128) -> Generator[FusionSegment, None, int]:
response = self._list_segments(offset, limit)
for item in response["segments"]:
segment = FusionSegment._from_client( # pylint: disable=protected-access
FusionSegmentClient(item["name"], self))
segment.description = item["description"]
yield segment
return response["totalCount"] # type: ignore[no-any-return]
def get_or_create_segment(self,
name: str = "default") -> FusionSegmentClient:
"""Get or create a fusion segment with the given name.
Arguments:
name: The name of the fusion segment.
Returns:
The created :class:`~tensorbay.client.segment.FusionSegmentClient` with given name.
"""
self._status.check_authority_for_draft()
if name not in self.list_segment_names():
self._create_segment(name)
return FusionSegmentClient(name, self)
def get_segment(self, name: str = "default") -> FusionSegmentClient:
"""Get a fusion segment in a certain commit according to given name.
Arguments:
name: The name of the required fusion segment.
Returns:
The required :class:`~tensorbay.client.segment.FusionSegmentClient`.
Raises:
ResourceNotExistError: When the required fusion segment does not exist.
"""
if name not in self.list_segment_names():
raise ResourceNotExistError(resource="segment",
identification=name)
return FusionSegmentClient(name, self)
def create_segment(self, name: str = "default") -> FusionSegmentClient:
"""Create a fusion segment with the given name.
Arguments:
name: The name of the fusion segment.
Returns:
The created :class:`~tensorbay.client.segment.FusionSegmentClient` with given name.
Raises:
NameConflictError: When the segment exists.
"""
self._status.check_authority_for_draft()
if name not in self.list_segment_names():
self._create_segment(name)
else:
raise NameConflictError(resource="segment", identification=name)
return FusionSegmentClient(name, self)
def copy_segment(
self,
source_name: str,
target_name: Optional[str] = None,
*,
source_client: Optional["FusionDatasetClient"] = None,
strategy: str = "abort",
) -> FusionSegmentClient:
"""Copy segment to this dataset.
Arguments:
source_name: The source name of the copied segment.
target_name: The target name of the copied segment.
This argument is used to specify a new name of the copied segment.
If None, the name of the copied segment will not be changed after copy.
source_client: The source dataset client of the copied segment.
This argument is used to specify where the copied segment comes from when the
copied segment is from another commit, draft or even another dataset.
If None, the copied segment comes from this dataset.
strategy: The strategy of handling the name conflict. There are three options:
1. "abort": stop copying and raise exception;
2. "override": the source segment will override the origin segment;
3. "skip": keep the origin segment.
Returns:
The client of the copied target segment.
"""
if not target_name:
target_name = source_name
self._copy_segment(source_name,
target_name,
source_client=source_client,
strategy=strategy)
return FusionSegmentClient(target_name, self)
def move_segment(
self,
source_name: str,
target_name: str,
*,
strategy: str = "abort",
) -> FusionSegmentClient:
"""Move/Rename segment in this dataset.
Arguments:
source_name: The source name of the moved segment.
target_name: The target name of the moved segment.
strategy: The strategy of handling the name conflict. There are three options:
1. "abort": stop moving and raise exception;
2. "override": the source segment will override the origin segment;
3. "skip": keep the origin segment.
Returns:
The client of the moved target segment.
"""
self._move_segment(source_name, target_name, strategy=strategy)
return FusionSegmentClient(target_name, self)
def test__upload_segment(self, mocker):
segment_test = FusionSegment(name="test1")
ulids = []
done_frames = []
for i in range(5):
temp_frame = Frame()
temp_ulid = from_timestamp(10 * i + 10)
temp_frame.frame_id = temp_ulid
if i % 2 == 0:
temp_frame["camera"] = Data(f"{i}.png")
done_frames.append(temp_frame)
else:
temp_frame["lidar"] = Data(f"{i}.png")
ulids.append(temp_ulid)
segment_test.append(temp_frame)
segment_client = FusionSegmentClient(
name="test1", data_client=self.fusion_dataset_client)
get_or_create_segment = mocker.patch(
f"{dataset.__name__}.FusionDatasetClient.get_or_create_segment",
return_value=segment_client,
)
list_frames = mocker.patch(
f"{segment.__name__}.FusionSegmentClient.list_frames",
return_value=done_frames,
)
multithread_upload = mocker.patch(
f"{dataset.__name__}.multithread_upload")
with Tqdm(5, disable=False) as pbar:
self.fusion_dataset_client._upload_segment(
segment_test, jobs=8, skip_uploaded_files=True, pbar=pbar)
get_or_create_segment.assert_called_once_with(segment_test.name)
list_frames.assert_called_once_with()
args, keywords = multithread_upload.call_args
for index, values in enumerate(args[1]):
data, sensor_name, frame_id = values
assert data.path == f"{index * 2 + 1}.png"
assert sensor_name == "lidar"
assert frame_id == ulids[index * 2 + 1].str
assert keywords[
"callback"] == segment_client._synchronize_upload_info
assert keywords["jobs"] == 8
assert keywords["pbar"] == pbar
multithread_upload.assert_called_once()
with Tqdm(5, disable=False) as pbar:
self.fusion_dataset_client._upload_segment(
segment_test, jobs=8, skip_uploaded_files=False, pbar=pbar)
get_or_create_segment.assert_called_with(segment_test.name)
list_frames.assert_called_with()
args, keywords = multithread_upload.call_args
for index, values in enumerate(args[1]):
data, sensor_name, frame_id = values
assert data.path == f"{index}.png"
if index % 2 == 0:
assert sensor_name == "camera"
else:
assert sensor_name == "lidar"
assert frame_id == ulids[index].str
assert keywords[
"callback"] == segment_client._synchronize_upload_info
assert keywords["jobs"] == 8
assert keywords["pbar"] == pbar