diff --git a/backend/labels/managers.py b/backend/labels/managers.py index f81cfcdb..fb4cda23 100644 --- a/backend/labels/managers.py +++ b/backend/labels/managers.py @@ -81,3 +81,8 @@ class RelationManager(LabelManager): def can_annotate(self, label, project) -> bool: return True + + +class BoundingBoxManager(LabelManager): + def can_annotate(self, label, project) -> bool: + return True diff --git a/backend/labels/migrations/0015_create_boundingbox_table.py b/backend/labels/migrations/0015_create_boundingbox_table.py new file mode 100644 index 00000000..ea861d41 --- /dev/null +++ b/backend/labels/migrations/0015_create_boundingbox_table.py @@ -0,0 +1,61 @@ +# Generated by Django 4.0.4 on 2022-06-29 06:19 + +from django.conf import settings +from django.db import migrations, models +import django.db.models.deletion +import uuid + + +class Migration(migrations.Migration): + + dependencies = [ + ("label_types", "0007_delete_relationtypeold"), + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ("examples", "0006_alter_example_upload_name"), + ("labels", "0014_remove_uuid_null"), + ] + + operations = [ + migrations.CreateModel( + name="BoundingBox", + fields=[ + ("id", models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")), + ("uuid", models.UUIDField(default=uuid.uuid4, unique=True)), + ("prob", models.FloatField(default=0.0)), + ("manual", models.BooleanField(default=False)), + ("created_at", models.DateTimeField(auto_now_add=True)), + ("updated_at", models.DateTimeField(auto_now=True)), + ("x", models.FloatField()), + ("y", models.FloatField()), + ("width", models.FloatField()), + ("height", models.FloatField()), + ( + "example", + models.ForeignKey( + on_delete=django.db.models.deletion.CASCADE, related_name="bboxes", to="examples.example" + ), + ), + ( + "label", + models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to="label_types.categorytype"), + ), + ("user", models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL)), + ], + ), + migrations.AddConstraint( + model_name="boundingbox", + constraint=models.CheckConstraint(check=models.Q(("x__gte", 0)), name="x >= 0"), + ), + migrations.AddConstraint( + model_name="boundingbox", + constraint=models.CheckConstraint(check=models.Q(("y__gte", 0)), name="y >= 0"), + ), + migrations.AddConstraint( + model_name="boundingbox", + constraint=models.CheckConstraint(check=models.Q(("width__gte", 0)), name="width >= 0"), + ), + migrations.AddConstraint( + model_name="boundingbox", + constraint=models.CheckConstraint(check=models.Q(("height__gte", 0)), name="height >= 0"), + ), + ] diff --git a/backend/labels/models.py b/backend/labels/models.py index 0c05f176..548be080 100644 --- a/backend/labels/models.py +++ b/backend/labels/models.py @@ -5,6 +5,7 @@ from django.core.exceptions import ValidationError from django.db import models from .managers import ( + BoundingBoxManager, CategoryManager, LabelManager, RelationManager, @@ -126,3 +127,21 @@ class Relation(Label): if not same_example: raise ValidationError("You need to label the same example.") return super().clean() + + +class BoundingBox(Label): + objects = BoundingBoxManager() + x = models.FloatField() + y = models.FloatField() + width = models.FloatField() + height = models.FloatField() + label = models.ForeignKey(to=CategoryType, on_delete=models.CASCADE) + example = models.ForeignKey(to=Example, on_delete=models.CASCADE, related_name="bboxes") + + class Meta: + constraints = [ + models.CheckConstraint(check=models.Q(x__gte=0), name="x >= 0"), + models.CheckConstraint(check=models.Q(y__gte=0), name="y >= 0"), + models.CheckConstraint(check=models.Q(width__gte=0), name="width >= 0"), + models.CheckConstraint(check=models.Q(height__gte=0), name="height >= 0"), + ] diff --git a/backend/labels/tests/test_bbox.py b/backend/labels/tests/test_bbox.py new file mode 100644 index 00000000..7a014ecd --- /dev/null +++ b/backend/labels/tests/test_bbox.py @@ -0,0 +1,29 @@ +from django.db import IntegrityError +from django.test import TestCase +from model_mommy import mommy + +from projects.tests.utils import prepare_project + + +class TestBoundingBox(TestCase): + @classmethod + def setUpTestData(cls): + cls.project = prepare_project() + cls.example = mommy.make("Example", project=cls.project.item) + cls.user = cls.project.admin + + def test_cannot_create_label_if_x_is_less_than_zero(self): + with self.assertRaises(IntegrityError): + mommy.make("BoundingBox", example=self.example, x=-1, y=0, width=0, height=0) + + def test_cannot_create_label_if_y_is_less_than_zero(self): + with self.assertRaises(IntegrityError): + mommy.make("BoundingBox", example=self.example, x=0, y=-1, width=0, height=0) + + def test_cannot_create_label_if_width_is_less_than_zero(self): + with self.assertRaises(IntegrityError): + mommy.make("BoundingBox", example=self.example, x=0, y=0, width=-1, height=0) + + def test_cannot_create_label_if_height_is_less_than_zero(self): + with self.assertRaises(IntegrityError): + mommy.make("BoundingBox", example=self.example, x=0, y=0, width=0, height=-1)