def test_delete_task_failed(self):
resp = Response()
resp.status_code = 404
self.session_mock.delete.return_value = resp
with self.assertRaises(CVATServiceException):
CvatService().delete_task(task_id=1)
def test_get_users_request_fails(self):
self.session_mock.post.return_value.raise_for_status.side_effect = HTTPError
with self.assertRaisesRegex(
CVATServiceException,
"Request to 'http://localhost:8080/api/v1/auth/login' failed"):
CvatService().get_users()
def test_delete_task(self):
resp = Response()
resp.status_code = 204
self.session_mock.delete.return_value = resp
CvatService().delete_task(task_id=1)
self.session_mock.delete.assert_called_once()
def get(self, request, *args, **kwargs):
if request.GET.get('cvat', 'false').lower() == 'true':
try:
CvatService()
except Exception as e:
return HttpResponse(
content=f"Failed to connect to CVAT: {e}",
status=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
return HttpResponse("ok")
def test_get_root_user(self):
user = CvatService().get_root_user()
self.assertEqual(
user,
{
"url": "http://localhost:8080/api/v1/users/2",
"id": 2,
"username": "******",
"first_name": "",
"last_name": "",
"email": "",
},
)
def test_get_annotations_not_created(self, sleep_mock):
self.session_mock.get.return_value = mock.Mock(
status_code=202,
content=b'',
)
with self.assertRaisesRegex(CVATServiceException,
r"Failed to get annotations \[202\]: b''"):
CvatService().get_annotations(
task_id=1,
task_name='test',
)
self.assertEqual(sleep_mock.call_count, 10)
def test_create_task(self):
data = {
'url': 'http://localhost:8080/api/v1/tasks/1',
'id': 1,
'name': 'test',
'size': 0,
'mode': '',
'owner': 1,
'assignee': 2,
'bug_tracker': '',
'created_date': '2020-05-07T18:17:39.484093Z',
'updated_date': '2020-05-07T18:17:39.484127Z',
'overlap': None,
'segment_size': 10,
'z_order': False,
'status': 'annotation',
'labels': [
{
'id': 1,
'name': 'newLabel',
'attributes': [],
},
],
'segments': [],
'image_quality': 70,
'start_frame': 0,
'stop_frame': 0,
'frame_filter': 'step=1',
'project': None,
}
self.session_mock.post.return_value = mock.Mock(
status_code=201,
content=dumps(data),
json=mock.Mock(return_value=data, ),
)
task = CvatService().create_task(
name='test',
assignee_id=2,
owner_id=1,
remote_files=[
"https://d3o54g14k1n39o.cloudfront.net/batch_2020-05-08/image.jpg",
],
)
self.assertEqual(task['id'], 1)
self.assertEqual(task['name'], 'test')
self.assertEqual(task['assignee'], 2)
self.assertEqual(task['owner'], 1)
def test_get_users(self):
users = CvatService().get_users()
self.assertEqual(len(users), 2)
self.assertEqual(
users[0],
{
"url": "http://*****:*****@labler.com",
},
)
def test_get_task(self):
self.session_mock.get.return_value.json.return_value = {
'id': 1,
'name': 'foo',
}
response = CvatService().get_task(task_id=1)
self.assertEqual(
response,
{
'id': 1,
'name': 'foo',
},
)
def test_get_user(self):
self.session_mock.get.return_value.json.return_value = {
'url': 'http://*****:*****@labler.com',
'groups': [
'annotator',
],
'is_staff': False,
'is_superuser': False,
'is_active': True,
'last_login': '******',
'date_joined': '2020-05-04T17:03:01Z',
}
user = CvatService().get_user(user_id=1)
self.assertEqual(user['username'], 'labler')
def test_get_annotations(self, sleep_mock):
self.session_mock.get.side_effect = [
mock.Mock(
status_code=202,
content=b'',
),
mock.Mock(
status_code=201,
content=b'',
),
mock.Mock(
status_code=200,
content=b'content',
),
]
response = CvatService().get_annotations(
task_id=1,
task_name='test',
)
self.assertEqual(response, b'content')
self.assertEqual(sleep_mock.call_count, 2)
def test_get_root_user_not_found(self):
self.session_mock.get.return_value.json.return_value = {'results': []}
with self.assertRaisesRegex(RuntimeError,
"Failed to find root CVAT user: admin"):
CvatService().get_root_user()
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._cvat_service = CvatService()
class Command(BaseCommand):
help = "Process tasks statuses"
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._cvat_service = CvatService()
def handle(self, *args, **kwargs):
try:
labeling_tasks = list(
LabelingTask.objects.filter(
status__in=(
LabelingTaskStatus.ANNOTATION,
LabelingTaskStatus.VALIDATION,
LabelingTaskStatus.COMPLETED,
),
error__isnull=True,
).order_by("updated_at")
[:settings.TASK_STATUSES_WORKER_CHUNK_SIZE], )
if labeling_tasks:
self._process_labeling_tasks(labeling_tasks=labeling_tasks)
else:
logger.info("No pending labeling tasks found")
except Exception as e:
raise CommandError(f"Failed to process labeling tasks: {e}")
def _process_labeling_tasks(self,
labeling_tasks: List[LabelingTask]) -> None:
logger.info(f"Processing {len(labeling_tasks)} labeling tasks")
# NOTE: Make sure we get the "oldest" tasks every time, so we need to update
# updated_at to the current time for the current chunk.
LabelingTask.objects.filter(
pk__in=[task.pk for task in labeling_tasks]).update(
updated_at=timezone.now())
# update labeling tasks statuses
labeling_tasks_to_upload = []
for labeling_task, result_future in self._get_cvat_statuses(
labeling_tasks=labeling_tasks):
try:
cvat_task = result_future.result()
except Exception as e:
labeling_task.set_failed(
error=f"Failed to get task status: {e}")
else:
if labeling_task.status in (LabelingTaskStatus.ANNOTATION,
LabelingTaskStatus.VALIDATION):
cvat_task_status = cvat_task.get('status')
if cvat_task_status is not None and labeling_task.status != cvat_task_status:
labeling_task.update_status(status=cvat_task_status)
if labeling_task.status == LabelingTaskStatus.COMPLETED:
labeling_tasks_to_upload.append(labeling_task)
# upload annotations
for labeling_task, result_future in self._upload_annotations(
labeling_tasks=labeling_tasks_to_upload):
try:
result_future.result()
except NoAnnotationException as noAnnotationException:
logger.error(f"{noAnnotationException}")
labeling_task.set_failed(error=f"{noAnnotationException}")
except Exception as e:
logger.error(f"{e}")
labeling_task.set_failed(error=f"{e}")
else:
labeling_task.update_status(status=LabelingTaskStatus.SAVED)
def _get_cvat_statuses(self, labeling_tasks: List[LabelingTask]):
with ThreadPoolExecutor() as executor:
future_to_labeling_task = {
executor.submit(self._cvat_service.get_task, t.task_id): t
for t in labeling_tasks
}
for future in as_completed(future_to_labeling_task):
labeling_task = future_to_labeling_task[future]
yield labeling_task, future
def _upload_annotations(self, labeling_tasks: List[LabelingTask]):
logger.info(
f"Uploading annotations for {len(labeling_tasks)} labeling tasks")
with ThreadPoolExecutor() as executor:
future_to_labeling_task = {
executor.submit(self._upload_labeling_task_annotations, t): t
for t in labeling_tasks
}
for future in as_completed(future_to_labeling_task):
labeling_task = future_to_labeling_task[future]
yield labeling_task, future
@staticmethod
def _get_cvat_zip_folderpath(annotations_format: str):
if annotations_format == AnnotationsFormat.YOLO_ZIP_1_1:
return 'obj_train_data'
else:
return 'Annotations'
@staticmethod
def _get_label_file_extension(annotations_format: str):
if annotations_format == AnnotationsFormat.YOLO_ZIP_1_1:
return '.txt'
else:
return '.xml'
@staticmethod
def _get_annotations_format(dataset: Dataset):
if dataset.dataset_format == DatasetFormat.YOLO:
return AnnotationsFormat.YOLO_ZIP_1_1
else:
return AnnotationsFormat.PASCAL_VOB_ZIP_1_1
def _upload_labeling_task_annotations(self, labeling_task: LabelingTask):
dataset = labeling_task.media_assets.first().dataset
annotation_frmt = self._get_annotations_format(dataset)
try:
annotations_content_zip = self._cvat_service.get_annotations(
task_id=labeling_task.task_id,
task_name=labeling_task.name,
annotation_format=annotation_frmt,
)
except Exception as e:
raise Exception(f"Failed to get task annotations: {e}")
zip_fp = BytesIO(annotations_content_zip)
zf = ZipFile(file=zip_fp)
annotation_filenames = {
os.path.split(zi.filename)[-1]: zf.open(zi)
for zi in zf.filelist if zi.filename.startswith(
self._get_cvat_zip_folderpath(annotation_frmt))
}
media_assets = labeling_task.media_assets.all()
media_assets_filenames = {
f"{os.path.splitext(media_asset.filename)[0]}" +
self._get_label_file_extension(annotation_frmt)
for media_asset in media_assets
}
logger.info(media_assets_filenames)
logger.info(annotation_filenames)
if len(annotation_filenames) == 0:
raise NoAnnotationException(
f"Missing all task annotations for task :{str(labeling_task.task_id)}"
)
for media_asset in media_assets:
try:
asset_filename = os.path.splitext(media_asset.filename)[0]
bucket_name = media_asset.dataset.bucket.name
s3_client = S3Client(bucket_name=bucket_name)
labeling_file_name = f"{asset_filename}" + self._get_label_file_extension(
annotation_frmt)
if labeling_file_name in annotation_filenames:
file_object = annotation_filenames[
f"{asset_filename}" +
self._get_label_file_extension(annotation_frmt)]
# TODO:remove deprecated remote_label_path property
media_asset.labeling_asset_filepath = media_asset.remote_label_path
s3_client.upload_file_obj(
file_obj=file_object,
bucket=bucket_name,
key=media_asset.full_label_path,
)
media_asset.save()
logger.info(
f"Uploaded annotation '{media_asset.full_label_path}'")
except Exception as e:
raise Exception(f"Failed to upload task annotations: {e}")