You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

107 lines
2.7 KiB

  1. from typing import Any, Dict, Protocol, Tuple
  2. from django.db import models
  3. from examples.models import Example
  4. from labels.models import BoundingBox, Category, Relation, Segmentation, Span, TextLabel
  5. from projects.models import Project
  6. DATA = "data"
  7. class ExportedExampleManager(models.Manager):
  8. def confirmed(self, project: Project, user=None):
  9. if project.collaborative_annotation:
  10. return self.filter(project=project).exclude(states=None)
  11. else:
  12. assert user is not None
  13. return self.filter(project=project, states__confirmed_by=user)
  14. class ExportedExample(Example):
  15. objects = ExportedExampleManager()
  16. def to_dict(self, is_text_project=True) -> Dict[str, Any]:
  17. return {"id": self.id, DATA: self.text if is_text_project else self.upload_name, **self.meta}
  18. class Meta:
  19. proxy = True
  20. class ExportedLabel(Protocol):
  21. objects: models.Manager
  22. def to_dict(self) -> Dict[str, Any]:
  23. raise NotImplementedError("Please implement this method in the subclass.")
  24. def to_string(self) -> str:
  25. raise NotImplementedError("Please implement this method in the subclass.")
  26. def to_tuple(self) -> Tuple:
  27. raise NotImplementedError("Please implement this method in the subclass.")
  28. class ExportedCategory(Category):
  29. def to_string(self) -> str:
  30. return self.label.text
  31. class Meta:
  32. proxy = True
  33. class ExportedSpan(Span):
  34. def to_dict(self):
  35. return {
  36. "id": self.id,
  37. "label": self.label.text,
  38. "start_offset": self.start_offset,
  39. "end_offset": self.end_offset,
  40. }
  41. def to_tuple(self):
  42. return self.start_offset, self.end_offset, self.label.text
  43. class Meta:
  44. proxy = True
  45. class ExportedRelation(Relation):
  46. def to_dict(self):
  47. return {"id": self.id, "from_id": self.from_id.id, "to_id": self.to_id.id, "type": self.type.text}
  48. class Meta:
  49. proxy = True
  50. class ExportedText(TextLabel):
  51. def to_string(self) -> str:
  52. return self.text
  53. class Meta:
  54. proxy = True
  55. class ExportedBoundingBox(BoundingBox):
  56. def to_dict(self):
  57. return {
  58. "uuid": str(self.uuid),
  59. "x": self.x,
  60. "y": self.y,
  61. "width": self.width,
  62. "height": self.height,
  63. "label": self.label.text,
  64. }
  65. def to_tuple(self):
  66. return self.x, self.y, self.width, self.height
  67. class Meta:
  68. proxy = True
  69. class ExportedSegmentation(Segmentation):
  70. def to_dict(self):
  71. return {"uuid": str(self.uuid), "points": self.points, "label": self.label.text}
  72. class Meta:
  73. proxy = True