# Copyright (C) 2019-2022 Intel Corporation # Copyright (C) CVAT.ai Corporation # # SPDX-License-Identifier: MIT import io import itertools from collections import OrderedDict, defaultdict from collections.abc import Callable, Generator, Sequence from contextlib import nullcontext from copy import deepcopy from enum import Enum from datumaro.components.errors import DatasetError, DatasetImportError, DatasetNotFoundError from django.conf import settings from django.db import transaction from django.db.models.query import Prefetch, QuerySet from rest_framework.exceptions import ValidationError from cvat.apps.dataset_manager.annotation import AnnotationIR, AnnotationManager from cvat.apps.dataset_manager.bindings import ( CvatDatasetNotFoundError, CvatImportError, JobData, TaskData, ) from cvat.apps.dataset_manager.util import TmpDirManager, faster_deepcopy from cvat.apps.engine import models, serializers from cvat.apps.engine.log import DatasetLogManager from cvat.apps.engine.plugins import plugin_decorator from cvat.apps.engine.utils import av_scan_paths, take_by from cvat.apps.events.handlers import handle_annotations_change from cvat.apps.profiler import silk_profile from cvat.utils import django_database as db_utils dlogger = DatasetLogManager() class dotdict(OrderedDict): """dot.notation access to dictionary attributes""" __getattr__ = OrderedDict.get __setattr__ = OrderedDict.__setitem__ __delattr__ = OrderedDict.__delitem__ __eq__ = lambda self, other: self.id == other.id __hash__ = lambda self: self.id class PatchAction(str, Enum): CREATE = "create" UPDATE = "update" DELETE = "delete" @classmethod def values(cls): return [item.value for item in cls] def __str__(self): return self.value class AnnotationImportMode(str, Enum): REPLACE = "replace" APPEND = "append" @classmethod def values(cls): return [item.value for item in cls] def __str__(self): return self.value def _receive_attributes_from_db(related_manager, foreign_key: str) -> defaultdict[int, list]: attributes = defaultdict(list) for attr in related_manager.values( foreign_key, "spec_id", "value", "id", ).iterator(chunk_size=settings.DEFAULT_DB_ANNO_CHUNK_SIZE): attributes[attr[foreign_key]].append( dotdict( { "spec_id": attr["spec_id"], "value": attr["value"], "id": attr["id"], } ) ) return attributes def merge_table_rows(rows, keys_for_merge, field_id): # It is necessary to keep a stable order of original rows # (e.g. for tracked boxes). Otherwise prev_box.frame can be bigger # than next_box.frame. merged_rows = OrderedDict() # Group all rows by field_id. In grouped rows replace fields in # accordance with keys_for_merge structure. for row in rows: row_id = row[field_id] if not row_id in merged_rows: merged_rows[row_id] = dotdict(row) for key in keys_for_merge: merged_rows[row_id][key] = [] for key in keys_for_merge: item = dotdict({v.split("__", 1)[-1]: row[v] for v in keys_for_merge[key]}) if item.id is not None: merged_rows[row_id][key].append(item) # Remove redundant keys from final objects redundant_keys = [item for values in keys_for_merge.values() for item in values] for i in merged_rows: for j in redundant_keys: del merged_rows[i][j] return list(merged_rows.values()) def _validate_input_annotations( annotations: AnnotationIR | dict, *, db_data: models.Data, dimension: models.DimensionType ) -> AnnotationIR: if not isinstance(annotations, AnnotationIR): annotations = AnnotationIR(dimension, annotations) if annotations.tracks and db_data.validation_mode == models.ValidationMode.GT_POOL: # Only tags and shapes can be used in tasks with GT pool raise ValidationError( "Tracks are not supported when task validation mode is {}".format( models.ValidationMode.GT_POOL ) ) if annotations.intervals: task_start = db_data.start_frame task_stop = db_data.stop_frame for interval in annotations.intervals: if interval["stop"] is not None and interval["start"] > interval["stop"]: interval_ref = ( f"Interval {interval['id']}" if interval.get("id") is not None else "Interval" ) raise ValidationError( f"{interval_ref} start must be <= stop, got " f"[{interval['start']}, {interval['stop']}]" ) if not annotations.is_interval_inside(interval, task_start, task_stop): interval_ref = ( f"Interval {interval['id']}" if interval.get("id") is not None else "Interval" ) raise ValidationError( f"{interval_ref} cannot be outside the task boundaries" f"[{task_start}, {task_stop}], got " f"[{interval['start']}, {interval['stop']}]" ) return annotations class JobAnnotation: @classmethod def add_prefetch_info( cls, queryset: QuerySet[models.Job], prefetch_images: bool = True ) -> QuerySet[models.Job]: assert issubclass(queryset.model, models.Job) label_qs = db_utils.add_prefetch_fields( models.Label.objects.all(), [ "skeleton", "parent", "attributespec_set", ], ) label_qs = JobData.add_prefetch_info(label_qs) task_data_queryset = models.Data.objects.all() if prefetch_images: task_data_queryset = task_data_queryset.select_related("video").prefetch_related( Prefetch("images", queryset=models.Image.objects.order_by("frame")) ) return queryset.select_related( "segment", "segment__task", ).prefetch_related( "segment__task__project", "segment__task__owner", "segment__task__assignee", Prefetch("segment__task__data", queryset=task_data_queryset), Prefetch("segment__task__label_set", queryset=label_qs), Prefetch("segment__task__project__label_set", queryset=label_qs), ) def __init__( self, pk, *, lock_job_in_db: bool = False, queryset: QuerySet | None = None, prefetch_images: bool = False, db_job: models.Job | None = None, ): assert db_job is None or lock_job_in_db is False assert (db_job is None and queryset is None) or prefetch_images is False assert db_job is None or queryset is None if db_job is None: if queryset is None: queryset = self.add_prefetch_info( models.Job.objects, prefetch_images=prefetch_images ) if lock_job_in_db: queryset = queryset.select_for_update() self.db_job: models.Job = db_utils.get_cached(queryset, pk=int(pk)) else: self.db_job: models.Job = db_job db_segment = self.db_job.segment self.start_frame = db_segment.start_frame self.stop_frame = db_segment.stop_frame self.ir_data = AnnotationIR(db_segment.task.dimension) self.db_labels = { db_label.id: db_label for db_label in ( db_segment.task.project.label_set.all() if db_segment.task.project_id else db_segment.task.label_set.all() ) } self.db_attributes = {} for db_label in self.db_labels.values(): self.db_attributes[db_label.id] = { "mutable": OrderedDict(), "immutable": OrderedDict(), "all": OrderedDict(), } for db_attr in db_label.attributespec_set.all(): default_value = dotdict( [ ("spec_id", db_attr.id), ("value", db_attr.default_value), ] ) if db_attr.mutable: self.db_attributes[db_label.id]["mutable"][db_attr.id] = default_value else: self.db_attributes[db_label.id]["immutable"][db_attr.id] = default_value self.db_attributes[db_label.id]["all"][db_attr.id] = default_value def reset(self): self.ir_data.reset() def _validate_attribute_for_existence(self, db_attr_val, label_id, attr_type): if db_attr_val.spec_id not in self.db_attributes[label_id][attr_type]: raise ValidationError("spec_id `{}` is invalid".format(db_attr_val.spec_id)) def _validate_label_for_existence(self, label_id): if label_id not in self.db_labels: raise ValidationError("label_id `{}` is invalid".format(label_id)) def _add_missing_shape(self, track, first_shape): if first_shape["type"] == "skeleton": # in case with skeleton track we always expect to see one shape in track first_shape["frame"] = track["frame"] else: missing_shape = deepcopy(first_shape) missing_shape["frame"] = track["frame"] missing_shape["outside"] = True missing_shape.pop("id", None) track["shapes"].append(missing_shape) def _correct_frame_of_tracked_shapes(self, track): shapes = sorted(track["shapes"], key=lambda a: a["frame"]) first_shape = shapes[0] if shapes else None if first_shape and track["frame"] < first_shape["frame"]: self._add_missing_shape(track, first_shape) elif first_shape and first_shape["frame"] < track["frame"]: track["frame"] = first_shape["frame"] def _sync_frames(self, tracks, parent_track): if not tracks: return min_frame = tracks[0]["frame"] for track in tracks: if parent_track and parent_track.frame < track["frame"]: track["frame"] = parent_track.frame # track and its first shape must have the same frame self._correct_frame_of_tracked_shapes(track) if track["frame"] < min_frame: min_frame = track["frame"] if not parent_track: return if min_frame < parent_track.frame: # parent track cannot have a frame greater than the frame of the child track parent_tracked_shape = parent_track.shapes.first() parent_track.frame = min_frame parent_tracked_shape.frame = min_frame parent_tracked_shape.save() parent_track.save() for track in tracks: if parent_track.frame < track["frame"]: track["frame"] = parent_track.frame self._correct_frame_of_tracked_shapes(track) def _save_tracks_to_db(self, tracks): def create_tracks(tracks, parent_track=None): db_tracks = [] db_track_attr_vals = [] db_shapes = [] db_shape_attr_vals = [] self._sync_frames(tracks, parent_track) tracks = [track for track in tracks if track["shapes"]] for track in tracks: track_attributes = track.pop("attributes", []) shapes = track.pop("shapes") elements = track.pop("elements", []) db_track = models.LabeledTrack(job=self.db_job, parent=parent_track, **track) self._validate_label_for_existence(db_track.label_id) for attr in track_attributes: db_attr_val = models.LabeledTrackAttributeVal( **attr, job_id=self.db_job.id, track_id=len(db_tracks) ) self._validate_attribute_for_existence( db_attr_val, db_track.label_id, "immutable" ) db_track_attr_vals.append(db_attr_val) for shape_idx, shape in enumerate(shapes): shape_attributes = shape.pop("attributes", []) db_shape = models.TrackedShape(**shape, track_id=len(db_tracks)) for attr in shape_attributes: db_attr_val = models.TrackedShapeAttributeVal( **attr, shape_id=len(db_shapes), job_id=self.db_job.id, ) self._validate_attribute_for_existence( db_attr_val, db_track.label_id, "mutable" ) db_shape_attr_vals.append(db_attr_val) db_shapes.append(db_shape) shape["attributes"] = shape_attributes db_tracks.append(db_track) track["attributes"] = track_attributes track["shapes"] = shapes if elements or parent_track is None: track["elements"] = elements db_tracks = db_utils.bulk_create(models.LabeledTrack, db_tracks) for db_attr_val in db_track_attr_vals: db_attr_val.track_id = db_tracks[db_attr_val.track_id].id db_utils.bulk_create(models.LabeledTrackAttributeVal, db_track_attr_vals) for db_shape in db_shapes: db_shape.track_id = db_tracks[db_shape.track_id].id db_shapes = db_utils.bulk_create(models.TrackedShape, db_shapes) for db_attr_val in db_shape_attr_vals: db_attr_val.shape_id = db_shapes[db_attr_val.shape_id].id db_utils.bulk_create(models.TrackedShapeAttributeVal, db_shape_attr_vals) shape_idx = 0 for track, db_track in zip(tracks, db_tracks): track["id"] = db_track.id for shape in track["shapes"]: shape["id"] = db_shapes[shape_idx].id shape_idx += 1 create_tracks(track.get("elements", []), db_track) create_tracks(tracks) self.ir_data.tracks = tracks def _save_shapes_to_db(self, shapes): def create_shapes(shapes, parent_shape=None): db_shapes = [] db_attr_vals = [] for shape in shapes: attributes = shape.pop("attributes", []) shape_elements = shape.pop("elements", []) # FIXME: need to clamp points (be sure that all of them inside the image) # Should we check here or implement a validator? db_shape = models.LabeledShape(job=self.db_job, parent=parent_shape, **shape) self._validate_label_for_existence(db_shape.label_id) for attr in attributes: db_attr_val = models.LabeledShapeAttributeVal( **attr, job_id=self.db_job.id, shape_id=len(db_shapes) ) self._validate_attribute_for_existence(db_attr_val, db_shape.label_id, "all") db_attr_vals.append(db_attr_val) db_shapes.append(db_shape) shape["attributes"] = attributes if shape_elements or parent_shape is None: shape["elements"] = shape_elements db_shapes = db_utils.bulk_create(models.LabeledShape, db_shapes) for db_attr_val in db_attr_vals: db_attr_val.shape_id = db_shapes[db_attr_val.shape_id].id db_utils.bulk_create(models.LabeledShapeAttributeVal, db_attr_vals) for shape, db_shape in zip(shapes, db_shapes): shape["id"] = db_shape.id create_shapes(shape.get("elements", []), db_shape) create_shapes(shapes) self.ir_data.shapes = shapes def _save_tags_to_db(self, tags): db_tags = [] db_attr_vals = [] for tag in tags: attributes = tag.pop("attributes", []) db_tag = models.LabeledImage(job=self.db_job, **tag) self._validate_label_for_existence(db_tag.label_id) for attr in attributes: db_attr_val = models.LabeledImageAttributeVal(**attr, job_id=self.db_job.id) self._validate_attribute_for_existence(db_attr_val, db_tag.label_id, "all") db_attr_val.tag_id = len(db_tags) db_attr_vals.append(db_attr_val) db_tags.append(db_tag) tag["attributes"] = attributes db_tags = db_utils.bulk_create(models.LabeledImage, db_tags) for db_attr_val in db_attr_vals: db_attr_val.image_id = db_tags[db_attr_val.tag_id].id db_utils.bulk_create(models.LabeledImageAttributeVal, db_attr_vals) for tag, db_tag in zip(tags, db_tags): tag["id"] = db_tag.id self.ir_data.tags = tags def _save_intervals_to_db(self, intervals: Sequence[dict]): def write_objects(intervals: Sequence[dict]): db_intervals = [] db_attr_vals = [] for interval in intervals: attributes = interval.pop("attributes", []) db_interval = models.LabeledInterval(job=self.db_job, **interval) self._validate_label_for_existence(db_interval.label_id) for attr in attributes: db_attr_val = models.LabeledIntervalAttributeVal( **attr, job_id=self.db_job.id, interval_id=len(db_intervals) ) self._validate_attribute_for_existence(db_attr_val, db_interval.label_id, "all") db_attr_vals.append(db_attr_val) db_intervals.append(db_interval) interval["attributes"] = attributes db_intervals = db_utils.bulk_create(models.LabeledInterval, db_intervals) for db_attr_val in db_attr_vals: db_attr_val.interval_id = db_intervals[db_attr_val.interval_id].id db_utils.bulk_create(models.LabeledIntervalAttributeVal, db_attr_vals) for interval, db_interval in zip(intervals, db_intervals): interval["id"] = db_interval.id write_objects(intervals) self.ir_data.intervals = intervals def _set_updated_date(self): db_task = self.db_job.segment.task with transaction.atomic(): self.db_job.touch() db_task.touch() if db_project := db_task.project: db_project.touch() @staticmethod def _data_is_empty(data): return not (data["tags"] or data["shapes"] or data["tracks"] or data["intervals"]) def _create(self, data): self.reset() self._save_tags_to_db(data["tags"]) self._save_shapes_to_db(data["shapes"]) self._save_tracks_to_db(data["tracks"]) self._save_intervals_to_db(data["intervals"]) def create(self, data): data = self._validate_input_annotations(data) self._create(data) handle_annotations_change(self.db_job, self.data, "create") if not self._data_is_empty(self.data): self._set_updated_date() def put(self, data): data = self._validate_input_annotations(data) deleted_data = self._delete() handle_annotations_change(self.db_job, deleted_data, "delete") deleted_data_is_empty = self._data_is_empty(deleted_data) self._create(data) handle_annotations_change(self.db_job, self.data, "create") if not deleted_data_is_empty or not self._data_is_empty(self.data): self._set_updated_date() def update(self, data): data = self._validate_input_annotations(data) # in case with "update" must be called prior any annotations in database changes # as this annotations are used to count removed/added shapes handle_annotations_change(self.db_job, data.data, "update") self._delete(data) self._create(data) if not self._data_is_empty(self.data): self._set_updated_date() def _validate_input_annotations(self, data: AnnotationIR | dict) -> AnnotationIR: db_task = self.db_job.segment.task return _validate_input_annotations( data, db_data=db_task.require_data(), dimension=db_task.dimension ) def _delete_job_labeledimages(self, ids__UNSAFE: list[int]) -> None: # ids__UNSAFE is a list, received from the user # we MUST filter it by job_id additionally before applying to any queries ids = self.db_job.labeledimage_set.filter(pk__in=ids__UNSAFE).values_list("id", flat=True) models.LabeledImageAttributeVal.objects.filter(image_id__in=ids).delete() self.db_job.labeledimage_set.filter(pk__in=ids).delete() def _delete_job_labeledshapes( self, ids__UNSAFE: list[int], *, is_subcall: bool = False ) -> None: # ids__UNSAFE is a list, received from the user # we MUST filter it by job_id additionally before applying to any queries if is_subcall: ids = ids__UNSAFE else: ids = self.db_job.labeledshape_set.filter(pk__in=ids__UNSAFE).values_list( "id", flat=True ) child_ids = self.db_job.labeledshape_set.filter(parent_id__in=ids).values_list( "id", flat=True ) if len(child_ids): self._delete_job_labeledshapes(child_ids, is_subcall=True) models.LabeledShapeAttributeVal.objects.filter(shape_id__in=ids).delete() self.db_job.labeledshape_set.filter(pk__in=ids).delete() def _delete_job_labeledtracks( self, ids__UNSAFE: list[int], *, is_subcall: bool = False ) -> None: # ids__UNSAFE is a list, received from the user # we MUST filter it by job_id additionally before applying to any queries if is_subcall: ids = ids__UNSAFE else: ids = self.db_job.labeledtrack_set.filter(pk__in=ids__UNSAFE).values_list( "id", flat=True ) child_ids = self.db_job.labeledtrack_set.filter(parent_id__in=ids).values_list( "id", flat=True ) if len(child_ids): self._delete_job_labeledtracks(child_ids, is_subcall=True) models.TrackedShapeAttributeVal.objects.filter(shape__track_id__in=ids).delete() models.LabeledTrackAttributeVal.objects.filter(track_id__in=ids).delete() self.db_job.labeledtrack_set.filter(pk__in=ids).delete() def _delete_job_intervals(self, ids__UNSAFE: list[int], *, is_subcall: bool = False) -> None: # ids__UNSAFE is a list, received from the user # we MUST filter it by job_id additionally before applying to any queries if is_subcall: ids = ids__UNSAFE else: ids = self.db_job.labeledinterval_set.filter(pk__in=ids__UNSAFE).values_list( "id", flat=True ) models.LabeledIntervalAttributeVal.objects.filter(interval_id__in=ids).delete() self.db_job.labeledinterval_set.filter(pk__in=ids).delete() def _delete(self, data=None): deleted_data = {} if data is None: self.init_from_db() deleted_data = self.data models.clear_annotations_in_jobs([self.db_job.id]) else: labeledimage_ids = [image["id"] for image in data["tags"]] labeledshape_ids = [shape["id"] for shape in data["shapes"]] labeledtrack_ids = [track["id"] for track in data["tracks"]] labeledinterval_ids = [interval["id"] for interval in data["intervals"]] for labeledimage_ids_chunk in take_by(labeledimage_ids, chunk_size=1000): self._delete_job_labeledimages(labeledimage_ids_chunk) for labeledshape_ids_chunk in take_by(labeledshape_ids, chunk_size=1000): self._delete_job_labeledshapes(labeledshape_ids_chunk) for labeledtrack_ids_chunk in take_by(labeledtrack_ids, chunk_size=1000): self._delete_job_labeledtracks(labeledtrack_ids_chunk) for labeledinterval_ids_chunk in take_by(labeledinterval_ids, chunk_size=1000): self._delete_job_intervals(labeledinterval_ids_chunk) deleted_data = { "version": self.ir_data.version, "tags": data["tags"], "shapes": data["shapes"], "tracks": data["tracks"], "intervals": data["intervals"], } self.reset() return deleted_data def delete(self, data=None): deleted_data = self._delete(data) if not self._data_is_empty(deleted_data): self._set_updated_date() handle_annotations_change(self.db_job, deleted_data, "delete") return deleted_data @staticmethod def _extend_attributes(attributeval_set, default_attribute_values): shape_attribute_specs_set = set(attr.spec_id for attr in attributeval_set) for db_attr in default_attribute_values: if db_attr.spec_id not in shape_attribute_specs_set: attributeval_set.append( dotdict( [ ("spec_id", db_attr.spec_id), ("value", db_attr.value), ] ) ) def _init_tags_from_db(self): db_tags = [ dotdict(row) for row in self.db_job.labeledimage_set.values( "id", "frame", "label_id", "group", "source", ) .order_by("frame") .iterator(chunk_size=settings.DEFAULT_DB_ANNO_CHUNK_SIZE) ] labeledimage_attributes = _receive_attributes_from_db( self.db_job.labeledimageattributeval_set, "image_id", ) for db_tag in db_tags: db_tag.attributes = labeledimage_attributes[db_tag.id] self._extend_attributes( db_tag.attributes, self.db_attributes[db_tag.label_id]["all"].values() ) serializer = serializers.LabeledImageSerializerFromDB(db_tags, many=True) self.ir_data.tags = serializer.data def _init_shapes_from_db(self, *, streaming: bool = False): db_shapes = ( dotdict(row) for row in self.db_job.labeledshape_set.values( "id", "label_id", "type", "frame", "group", "source", "score", "occluded", "outside", "z_order", "rotation", "points", "parent", ) .order_by("frame") .iterator(chunk_size=settings.DEFAULT_DB_ANNO_CHUNK_SIZE) ) labeledshape_attributes = _receive_attributes_from_db( self.db_job.labeledshapeattributeval_set, "shape_id", ) def yield_shapes_for_one_frame(shapes: dict, elements): for shape_id, shape_elements in elements.items(): shapes[shape_id].elements = shape_elements serializer = serializers.LabeledShapeSerializerFromDB(list(shapes.values()), many=True) yield from serializer.data shapes.clear() elements.clear() def generate_shapes(): shapes = {} elements = {} for db_shape in db_shapes: if shapes and next(iter(shapes.values())).frame != db_shape.frame: yield from yield_shapes_for_one_frame(shapes, elements) db_shape.attributes = labeledshape_attributes[db_shape.id] self._extend_attributes( db_shape.attributes, self.db_attributes[db_shape.label_id]["all"].values() ) if db_shape["type"] == str(models.ShapeType.SKELETON): # skeletons themselves should not have points as they consist of other elements # here we ensure that it was initialized correctly db_shape["points"] = [] if db_shape.parent is None: db_shape.elements = [] shapes[db_shape.id] = db_shape else: if db_shape.parent not in elements: elements[db_shape.parent] = [] elements[db_shape.parent].append(db_shape) yield from yield_shapes_for_one_frame(shapes, elements) if streaming: assert transaction.get_connection().in_atomic_block shapes = generate_shapes() # starting generation to initialise db-side cursor buffer = [] try: buffer.append(next(shapes)) except StopIteration: pass self.ir_data.shapes = itertools.chain(buffer, shapes) else: self.ir_data.shapes = list(generate_shapes()) def _init_tracks_from_db(self): db_tracks = [ dotdict(row) for row in self.db_job.labeledtrack_set.values( "id", "frame", "label_id", "group", "source", "parent", ) .order_by("id") .iterator(chunk_size=settings.DEFAULT_DB_ANNO_CHUNK_SIZE) ] if not db_tracks: self.ir_data.tracks = [] return tracks_by_id = {db_track.id: db_track for db_track in db_tracks} for db_track in db_tracks: db_track.shapes = [] for track_ids_chunk in take_by(sorted(tracks_by_id), 1000): db_shapes = ( models.TrackedShape.objects.filter(track_id__in=track_ids_chunk) .values( "track_id", "type", "occluded", "z_order", "rotation", "points", "id", "frame", "outside", ) .order_by("track_id", "frame") .iterator(chunk_size=settings.DEFAULT_DB_ANNO_CHUNK_SIZE) ) for db_shape in db_shapes: track_id = db_shape.pop("track_id") tracks_by_id[track_id].shapes.append(dotdict(db_shape)) labeledtrack_attributes = _receive_attributes_from_db( self.db_job.labeledtrackattributeval_set, "track_id", ) trackedshape_attributes = _receive_attributes_from_db( self.db_job.trackedshapeattributeval_set, "shape_id", ) tracks = {} elements = {} for db_track in db_tracks: if not db_track["shapes"]: continue # A result table can consist many equal rows for track/shape attributes # We need filter unique attributes manually db_track["attributes"] = list(set(labeledtrack_attributes[db_track["id"]])) self._extend_attributes( db_track.attributes, self.db_attributes[db_track.label_id]["immutable"].values() ) default_attribute_values = self.db_attributes[db_track.label_id]["mutable"].values() for db_shape in db_track["shapes"]: db_shape["attributes"] = list(set(trackedshape_attributes[db_shape["id"]])) # in case of trackedshapes need to interpolate attribute values and extend it # by previous shape attribute values (not default values) self._extend_attributes(db_shape["attributes"], default_attribute_values) if db_shape["type"] == str(models.ShapeType.SKELETON): # skeletons themselves should not have points as they consist of other elements # here we ensure that it was initialized correctly db_shape["points"] = [] default_attribute_values = db_shape["attributes"] if db_track.parent is None: db_track.elements = [] tracks[db_track.id] = db_track else: if db_track.parent not in elements: elements[db_track.parent] = [] elements[db_track.parent].append(db_track) for track_id, track_elements in elements.items(): tracks[track_id].elements = track_elements serializer = serializers.LabeledTrackSerializerFromDB(list(tracks.values()), many=True) self.ir_data.tracks = serializer.data def _init_intervals_from_db(self): db_intervals = ( dotdict(row) for row in self.db_job.labeledinterval_set.values( "id", "label_id", "start", "stop", "group", "source", "score", ) .order_by("start") .iterator(chunk_size=settings.DEFAULT_DB_ANNO_CHUNK_SIZE) ) db_attributes = _receive_attributes_from_db( self.db_job.labeledintervalattributeval_set, "interval_id", ) def serialize(annotations: dict) -> Generator[dict, None, None]: serializer = serializers.LabeledIntervalSerializerFromDB( list(annotations.values()), many=True ) yield from serializer.data annotations.clear() def generate_annotations(): annotations = {} for db_interval in db_intervals: db_interval.attributes = db_attributes[db_interval.id] self._extend_attributes( db_interval.attributes, self.db_attributes[db_interval.label_id]["all"].values() ) annotations[db_interval.id] = db_interval yield from serialize(annotations) self.ir_data.intervals = list(generate_annotations()) def _init_version_from_db(self): self.ir_data.version = 0 # FIXME: should be removed in the future def init_from_db(self, *, streaming: bool = False): self._init_tags_from_db() self._init_shapes_from_db(streaming=streaming) self._init_tracks_from_db() self._init_intervals_from_db() self._init_version_from_db() @property def data(self): return self.ir_data.data def export( self, dst_file: io.BufferedWriter, exporter: Callable[..., None], *, host: str = "", temp_dir: str | None = None, **options, ): job_data = JobData( annotation_ir=self.ir_data, db_job=self.db_job, host=host, ) with ( TmpDirManager.get_tmp_directory_for_export( instance_type=self.db_job.__class__.__name__, ) if not temp_dir else nullcontext(temp_dir) ) as temp_dir: exporter(dst_file, temp_dir, job_data, **options) def import_annotations( self, src_file, importer, *, import_mode: AnnotationImportMode = AnnotationImportMode.REPLACE, **options, ): job_data = JobData( annotation_ir=AnnotationIR(self.db_job.segment.task.dimension), db_job=self.db_job, create_callback=self.create, ) import_mode = AnnotationImportMode(import_mode) if import_mode == AnnotationImportMode.REPLACE: self.delete() elif import_mode == AnnotationImportMode.APPEND: pass else: assert False, f"Unknown annotation import mode: {import_mode}" with TmpDirManager.get_tmp_directory() as temp_dir: try: importer(src_file, temp_dir, job_data, **options) except (DatasetNotFoundError, CvatDatasetNotFoundError) as not_found: if settings.CVAT_LOG_IMPORT_ERRORS: dlogger.log_import_error( entity="job", entity_id=self.db_job.id, format_name=importer.DISPLAY_NAME, base_error=str(not_found), dir_path=temp_dir, ) raise not_found self.create(job_data.data.slice(self.start_frame, self.stop_frame).serialize()) class TaskAnnotation: def __init__(self, pk, *, write_only: bool = False): self.db_task = models.Task.objects.prefetch_related( Prefetch("data__images", queryset=models.Image.objects.order_by("frame")) ).get(id=pk) self._write_only = write_only # TODO: maybe include consensus jobs except for task export requested_job_types = [models.JobType.ANNOTATION] if self.db_task.data.validation_mode == models.ValidationMode.GT_POOL: requested_job_types.append(models.JobType.GROUND_TRUTH) self.db_jobs = ( JobAnnotation.add_prefetch_info(models.Job.objects, prefetch_images=False) .filter(segment__task_id=pk, type__in=requested_job_types) .order_by("id") ) if not write_only: self.ir_data = AnnotationIR(self.db_task.dimension) def reset(self): self.ir_data.reset() def _patch_data(self, data: AnnotationIR | dict, action: PatchAction | None): data = self._validate_input_annotations(data) if self.db_task.data.validation_mode == models.ValidationMode.GT_POOL: self._preprocess_input_annotations_for_gt_pool_task(data, action=action) splitted_data = {} jobs = {} for db_job in self.db_jobs: jid = db_job.id start = db_job.segment.start_frame stop = db_job.segment.stop_frame jobs[jid] = {"start": start, "stop": stop} splitted_data[jid] = (data.slice(start, stop), db_job) for jid, (job_data, db_job) in splitted_data.items(): data = AnnotationIR(self.db_task.dimension) if action is None: data.data = put_job_data(jid, job_data, db_job=db_job) else: data.data = patch_job_data(jid, job_data, action, db_job=db_job) if not self._write_only: if data.version > self.ir_data.version: self.ir_data.version = data.version self._merge_data(data, jobs[jid]["start"]) def _merge_data(self, data: AnnotationIR, start_frame: int): annotation_manager = AnnotationManager(self.ir_data, dimension=self.db_task.dimension) annotation_manager.merge(data, start_frame, overlap=self.db_task.overlap) def put(self, data): self._patch_data(data, None) def create(self, data): self._patch_data(data, PatchAction.CREATE) def _preprocess_input_annotations_for_gt_pool_task( self, data: AnnotationIR | dict, *, action: PatchAction | None ) -> AnnotationIR: if not isinstance(data, AnnotationIR): data = AnnotationIR(self.db_task.dimension, data) if data.tracks: # Only tags and shapes are supported in tasks with GT pool raise ValidationError( "Tracks are not supported when task validation mode is {}".format( models.ValidationMode.GT_POOL ) ) gt_job = self.db_task.gt_job if gt_job is None: raise AssertionError(f"Can't find GT job in the task {self.db_task.id}") db_data = self.db_task.require_data() frame_step = db_data.get_frame_step() def _to_rel_frame(abs_frame: int) -> int: return (abs_frame - db_data.start_frame) // frame_step # Copy GT pool annotations into other jobs, with replacement of any existing annotations gt_abs_frame_set = sorted(gt_job.segment.frame_set) task_gt_honeypots: dict[int, int] = {} # real_id -> [placeholder_id, ...] task_gt_frames: set[int] = set() for abs_frame, abs_real_frame in ( self.db_task.data.images.filter(is_placeholder=True, real_frame__in=gt_abs_frame_set) .values_list("frame", "real_frame") .iterator(chunk_size=1000) ): frame = _to_rel_frame(abs_frame) task_gt_frames.add(frame) task_gt_honeypots.setdefault(_to_rel_frame(abs_real_frame), []).append(frame) gt_pool_frames = tuple(map(_to_rel_frame, gt_abs_frame_set)) if sorted(gt_pool_frames) != list(range(min(gt_pool_frames), max(gt_pool_frames) + 1)): raise AssertionError("Expected a continuous GT pool frame set") # to be used in slice() gt_annotations = data.slice(min(gt_pool_frames), max(gt_pool_frames)) if action and not (gt_annotations.tags or gt_annotations.shapes or gt_annotations.tracks): return if not ( # put action is None or action == PatchAction.CREATE ): # allow validation frame editing only with full task updates raise ValidationError( "Annotations on validation frames can only be edited via task import or the GT job" ) task_annotation_manager = AnnotationManager(data, dimension=self.db_task.dimension) task_annotation_manager.clear_frames(task_gt_frames) for ann_type, gt_annotation in itertools.chain( zip(itertools.repeat("tag"), gt_annotations.tags), zip(itertools.repeat("shape"), gt_annotations.shapes), ): for honeypot_frame_id in task_gt_honeypots.get( gt_annotation["frame"], [] # some GT frames may be unused ): copied_annotation = faster_deepcopy(gt_annotation) copied_annotation["frame"] = honeypot_frame_id for ann in itertools.chain( [copied_annotation], copied_annotation.get("elements", []) ): ann.pop("id", None) if ann_type == "tag": data.add_tag(copied_annotation) elif ann_type == "shape": data.add_shape(copied_annotation) else: assert False return data def _validate_input_annotations(self, data: AnnotationIR | dict) -> AnnotationIR: return _validate_input_annotations( data, db_data=self.db_task.require_data(), dimension=self.db_task.dimension ) def update(self, data): self._patch_data(data, PatchAction.UPDATE) def delete(self, data=None): if data: self._patch_data(data, PatchAction.DELETE) else: for db_job in self.db_jobs: delete_job_data(db_job.id, db_job=db_job) def init_from_db(self, *, streaming: bool = False): self.reset() db_jobs = self.db_jobs if not streaming: db_jobs = db_jobs.select_for_update() for db_job in db_jobs: if db_job.type == models.JobType.GROUND_TRUTH and ( self.db_task.data.validation_mode != models.ValidationMode.GT_POOL ): continue annotation = JobAnnotation(db_job.id, db_job=db_job) annotation.init_from_db(streaming=streaming) if annotation.ir_data.version > self.ir_data.version: self.ir_data.version = annotation.ir_data.version self._merge_data(annotation.ir_data, start_frame=db_job.segment.start_frame) def export( self, dst_file: io.BufferedWriter, exporter: Callable[..., None], *, host: str = "", temp_dir: str | None = None, **options, ): task_data = TaskData( annotation_ir=self.ir_data, db_task=self.db_task, host=host, ) with ( TmpDirManager.get_tmp_directory_for_export( instance_type=self.db_task.__class__.__name__, ) if not temp_dir else nullcontext(temp_dir) ) as temp_dir: exporter(dst_file, temp_dir, task_data, **options) def import_annotations( self, src_file, importer, *, import_mode: AnnotationImportMode = AnnotationImportMode.REPLACE, **options, ): task_data = TaskData( annotation_ir=AnnotationIR(self.db_task.dimension), db_task=self.db_task, create_callback=self.create, ) import_mode = AnnotationImportMode(import_mode) if import_mode == AnnotationImportMode.REPLACE: self.delete() elif import_mode == AnnotationImportMode.APPEND: pass else: assert False, f"Unknown annotation import mode: {import_mode}" with TmpDirManager.get_tmp_directory() as temp_dir: try: importer(src_file, temp_dir, task_data, **options) except (DatasetNotFoundError, CvatDatasetNotFoundError) as not_found: if settings.CVAT_LOG_IMPORT_ERRORS: dlogger.log_import_error( entity="task", entity_id=self.db_task.id, format_name=importer.DISPLAY_NAME, base_error=str(not_found), dir_path=temp_dir, ) raise not_found self.create(task_data.data.serialize()) @property def data(self): return self.ir_data.data @silk_profile(name="GET job data") @transaction.atomic def get_job_data(pk, *, streaming: bool = False): annotation = JobAnnotation(pk) annotation.init_from_db(streaming=streaming) return annotation.data @silk_profile(name="POST job data") @transaction.atomic def put_job_data(pk, data: AnnotationIR | dict, *, db_job: models.Job | None = None): annotation = JobAnnotation(pk, db_job=db_job) annotation.put(data) return annotation.data @silk_profile(name="UPDATE job data") @plugin_decorator @transaction.atomic def patch_job_data( pk, data: AnnotationIR | dict, action: PatchAction, *, db_job: models.Job | None = None ): annotation = JobAnnotation(pk, db_job=db_job) if action == PatchAction.CREATE: annotation.create(data) elif action == PatchAction.UPDATE: annotation.update(data) elif action == PatchAction.DELETE: return annotation.delete(data) return annotation.data @silk_profile(name="DELETE job data") @transaction.atomic def delete_job_data(pk, *, db_job: models.Job | None = None): annotation = JobAnnotation(pk, db_job=db_job) annotation.delete() @db_utils.transaction_with_repeatable_read() def export_job( job_id: int, dst_file: str, *, format_name: str, server_url: str | None = None, save_images=False, temp_dir: str | None = None, ): from cvat.apps.dataset_manager.formats.registry import make_exporter job = JobAnnotation(job_id, prefetch_images=True) job.init_from_db(streaming=True) exporter = make_exporter(format_name) with open(dst_file, "wb") as f: job.export(f, exporter, host=server_url, save_images=save_images, temp_dir=temp_dir) @silk_profile(name="GET task data") @transaction.atomic def get_task_data(pk): annotation = TaskAnnotation(pk) annotation.init_from_db() return annotation.data @silk_profile(name="POST task data") @transaction.atomic def put_task_data(pk, data): annotation = TaskAnnotation(pk) annotation.put(data) return annotation.data @silk_profile(name="UPDATE task data") @transaction.atomic def patch_task_data(pk, data, action): annotation = TaskAnnotation(pk) if action == PatchAction.CREATE: annotation.create(data) elif action == PatchAction.UPDATE: annotation.update(data) elif action == PatchAction.DELETE: annotation.delete(data) return annotation.data @silk_profile(name="DELETE task data") @transaction.atomic def delete_task_data(pk): annotation = TaskAnnotation(pk) annotation.delete() @db_utils.transaction_with_repeatable_read() def export_task( task_id: int, dst_file: str, *, format_name: str, server_url: str | None = None, save_images: bool = False, temp_dir: str | None = None, ): from cvat.apps.dataset_manager.formats.registry import make_exporter task = TaskAnnotation(task_id) task.init_from_db(streaming=True) exporter = make_exporter(format_name) with open(dst_file, "wb") as f: task.export(f, exporter, host=server_url, save_images=save_images, temp_dir=temp_dir) @transaction.atomic def import_task_annotations( src_file, task_id, format_name, conv_mask_to_poly, *, import_mode: AnnotationImportMode = AnnotationImportMode.REPLACE, ): from cvat.apps.dataset_manager.formats.registry import make_importer av_scan_paths(src_file) task = TaskAnnotation(task_id, write_only=True) importer = make_importer(format_name) with open(src_file, "rb") as f: try: task.import_annotations( f, importer, conv_mask_to_poly=conv_mask_to_poly, import_mode=import_mode, ) except (DatasetError, DatasetImportError, DatasetNotFoundError) as ex: raise CvatImportError(str(ex)) @transaction.atomic def import_job_annotations( src_file, job_id, format_name, conv_mask_to_poly, *, import_mode: AnnotationImportMode = AnnotationImportMode.REPLACE, ): from cvat.apps.dataset_manager.formats.registry import make_importer av_scan_paths(src_file) job = JobAnnotation(job_id, prefetch_images=True) importer = make_importer(format_name) with open(src_file, "rb") as f: try: job.import_annotations( f, importer, conv_mask_to_poly=conv_mask_to_poly, import_mode=import_mode, ) except (DatasetError, DatasetImportError, DatasetNotFoundError) as ex: raise CvatImportError(str(ex))