From fe795dd0c42f343288977040685d9df57863b173 Mon Sep 17 00:00:00 2001
From: Hironsan <light.tree.1.13@gmail.com>
Date: Thu, 30 Jun 2022 14:25:27 +0900
Subject: [PATCH] Add segmentation model

---
 backend/labels/admin.py                       |  8 +++-
 backend/labels/managers.py                    |  5 ++
 .../labels/migrations/0016_segmentation.py    | 46 +++++++++++++++++++
 backend/labels/models.py                      |  8 ++++
 backend/labels/serializers.py                 | 22 ++++++++-
 5 files changed, 87 insertions(+), 2 deletions(-)
 create mode 100644 backend/labels/migrations/0016_segmentation.py

diff --git a/backend/labels/admin.py b/backend/labels/admin.py
index 9c65d0a3..1819746d 100644
--- a/backend/labels/admin.py
+++ b/backend/labels/admin.py
@@ -1,6 +1,6 @@
 from django.contrib import admin
 
-from .models import BoundingBox, Category, Span, TextLabel
+from .models import BoundingBox, Category, Segmentation, Span, TextLabel
 
 
 class SpanAdmin(admin.ModelAdmin):
@@ -23,7 +23,13 @@ class BoundingBoxAdmin(admin.ModelAdmin):
     ordering = ("example",)
 
 
+class SegmentationAdmin(admin.ModelAdmin):
+    list_display = ("example", "label", "user", "points")
+    ordering = ("example",)
+
+
 admin.site.register(Category, CategoryAdmin)
 admin.site.register(Span, SpanAdmin)
 admin.site.register(TextLabel, TextLabelAdmin)
 admin.site.register(BoundingBox, BoundingBoxAdmin)
+admin.site.register(Segmentation, SegmentationAdmin)
diff --git a/backend/labels/managers.py b/backend/labels/managers.py
index fb4cda23..6d7540c7 100644
--- a/backend/labels/managers.py
+++ b/backend/labels/managers.py
@@ -86,3 +86,8 @@ class RelationManager(LabelManager):
 class BoundingBoxManager(LabelManager):
     def can_annotate(self, label, project) -> bool:
         return True
+
+
+class SegmentationManager(LabelManager):
+    def can_annotate(self, label, project) -> bool:
+        return True
diff --git a/backend/labels/migrations/0016_segmentation.py b/backend/labels/migrations/0016_segmentation.py
new file mode 100644
index 00000000..ab1005e5
--- /dev/null
+++ b/backend/labels/migrations/0016_segmentation.py
@@ -0,0 +1,46 @@
+# Generated by Django 4.0.4 on 2022-06-30 05:24
+
+import uuid
+
+import django.db.models.deletion
+from django.conf import settings
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("examples", "0006_alter_example_upload_name"),
+        ("label_types", "0007_delete_relationtypeold"),
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ("labels", "0015_create_boundingbox_table"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="Segmentation",
+            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)),
+                ("points", models.JSONField(default=list)),
+                (
+                    "example",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE, related_name="segmentations", 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)),
+            ],
+            options={
+                "abstract": False,
+            },
+        ),
+    ]
diff --git a/backend/labels/models.py b/backend/labels/models.py
index 548be080..e93a54e1 100644
--- a/backend/labels/models.py
+++ b/backend/labels/models.py
@@ -9,6 +9,7 @@ from .managers import (
     CategoryManager,
     LabelManager,
     RelationManager,
+    SegmentationManager,
     SpanManager,
     TextLabelManager,
 )
@@ -145,3 +146,10 @@ class BoundingBox(Label):
             models.CheckConstraint(check=models.Q(width__gte=0), name="width >= 0"),
             models.CheckConstraint(check=models.Q(height__gte=0), name="height >= 0"),
         ]
+
+
+class Segmentation(Label):
+    objects = SegmentationManager()
+    points = models.JSONField(default=list)
+    label = models.ForeignKey(to=CategoryType, on_delete=models.CASCADE)
+    example = models.ForeignKey(to=Example, on_delete=models.CASCADE, related_name="segmentations")
diff --git a/backend/labels/serializers.py b/backend/labels/serializers.py
index 01dd4b75..c6c1ecec 100644
--- a/backend/labels/serializers.py
+++ b/backend/labels/serializers.py
@@ -1,6 +1,6 @@
 from rest_framework import serializers
 
-from .models import BoundingBox, Category, Relation, Span, TextLabel
+from .models import BoundingBox, Category, Relation, Segmentation, Span, TextLabel
 from examples.models import Example
 from label_types.models import CategoryType, RelationType, SpanType
 
@@ -91,3 +91,23 @@ class BoundingBoxSerializer(serializers.ModelSerializer):
             "height",
         )
         read_only_fields = ("user",)
+
+
+class SegmentationSerializer(serializers.ModelSerializer):
+    example = serializers.PrimaryKeyRelatedField(queryset=Example.objects.all())
+    label = serializers.PrimaryKeyRelatedField(queryset=CategoryType.objects.all())
+
+    class Meta:
+        model = Segmentation
+        fields = (
+            "id",
+            "uuid",
+            "prob",
+            "user",
+            "example",
+            "created_at",
+            "updated_at",
+            "label",
+            "points",
+        )
+        read_only_fields = ("user",)