Added train, val, test selctor

backend/src/models/signvideo.py

@@ -9,6 +9,7 @@ class SignVideo(SQLModelExtended, table=True):
     id: int = Field(primary_key=True)
 
     approved: bool = False
+    dataset: str = "train" # train, test, val
 
     # foreign key to sign
     sign_id: int = Field(default=None, foreign_key="sign.id")
@@ -22,4 +23,5 @@ class SignVideo(SQLModelExtended, table=True):
 
 class SignVideoOut(BaseModel):
     id: int
-    approved: bool
+    approved: bool
+    dataset: str

backend/src/routers/signvideo.py

@@ -157,6 +157,7 @@ async def sign_video(
 
 class SignVideoUpdate(BaseModel):
     approved: bool
+    dataset: str
 
 @router.patch("/{video_id}/", status_code=status.HTTP_200_OK)
 async def sign_video(
@@ -178,7 +179,12 @@ async def sign_video(
     if not sign_video:
         raise BaseException("Sign video not found")
 
+    # check if dataset is train, val or test
+    if update.dataset not in ["train", "val", "test"]:
+        raise BaseException("Dataset must be train, val or test")
+
     sign_video.approved = update.approved
+    sign_video.dataset = update.dataset
     await sign_video.save(session)