From 4ccd9a0b8b0893853d64b8e81fedf11de97570ec Mon Sep 17 00:00:00 2001 From: Hironsan Date: Sat, 15 May 2021 09:58:59 +0900 Subject: [PATCH] Update get_annotation_serializer --- backend/api/serializers.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/backend/api/serializers.py b/backend/api/serializers.py index 5316beb2..8086018e 100644 --- a/backend/api/serializers.py +++ b/backend/api/serializers.py @@ -6,12 +6,12 @@ from rest_framework import serializers from rest_framework.exceptions import ValidationError from rest_polymorphic.serializers import PolymorphicSerializer -from .models import (DOCUMENT_CLASSIFICATION, SEQ2SEQ, SEQUENCE_LABELING, - SPEECH2TEXT, AutoLabelingConfig, Category, Comment, - Example, ImageClassificationProject, Label, Project, Role, - RoleMapping, Seq2seqProject, SequenceLabelingProject, - Span, Speech2textProject, Tag, TextClassificationProject, - TextLabel) +from .models import (DOCUMENT_CLASSIFICATION, IMAGE_CLASSIFICATION, SEQ2SEQ, + SEQUENCE_LABELING, SPEECH2TEXT, AutoLabelingConfig, + Category, Comment, Example, ImageClassificationProject, + Label, Project, Role, RoleMapping, Seq2seqProject, + SequenceLabelingProject, Span, Speech2textProject, Tag, + TextClassificationProject, TextLabel) class UserSerializer(serializers.ModelSerializer): @@ -320,7 +320,8 @@ def get_annotation_serializer(task: str): DOCUMENT_CLASSIFICATION: CategorySerializer, SEQUENCE_LABELING: SpanSerializer, SEQ2SEQ: TextLabelSerializer, - SPEECH2TEXT: TextLabelSerializer + SPEECH2TEXT: TextLabelSerializer, + IMAGE_CLASSIFICATION: CategorySerializer, } try: return mapping[task]