app/routers/video.py

import asyncio
import os
import re
import shutil
import time
import aiofiles
import cv2

from multiprocessing import Process
from fastapi import (
    APIRouter,
    Depends,
    HTTPException,
    UploadFile,
    BackgroundTasks,
    status,
)
from firebase_admin import messaging
from app import db
from app import supabase
from app.dependencies import get_current_user
from app.routers.image import inferenceImage
from google.cloud.firestore_v1.base_query import FieldFilter

router = APIRouter(prefix="/video", tags=["Video"])


@router.post("")
async def handleVideoRequest(
    file: UploadFile,
    background_tasks: BackgroundTasks,
    threshold: float = 0.3,
    user=Depends(get_current_user),
):
    if re.search("^video\/", file.content_type) is None:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="File must be video",
        )

    try:
        if user["sub"] is None:
            return HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST, detail="User not found"
            )
        id = str(now())
        _, artifact_ref = db.collection("artifacts").add(
            {"name": id + ".mp4", "status": "pending"}
        )
        os.mkdir(id)
        async with aiofiles.open(os.path.join(id, "input.mp4"), "wb") as out_file:
            while content := await file.read(1024):
                await out_file.write(content)
        background_tasks.add_task(inferenceVideo, artifact_ref.id, id, threshold)
        return id + ".mp4"
    except ValueError as err:
        print(err)
        print("Error processing video")
        shutil.rmtree(id)


def now():
    return round(time.time() * 1000)


async def inferenceVideo(artifactId: str, inputDir: str, threshold: float):
    try:
        Process(updateArtifact(artifactId, {"status": "processing"})).start()
        cap = cv2.VideoCapture(
            filename=os.path.join(inputDir, "input.mp4"), apiPreference=cv2.CAP_FFMPEG
        )
        fps = cap.get(cv2.CAP_PROP_FPS)
        size = (
            int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
            int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
        )
        result = cv2.VideoWriter(
            filename=os.path.join(inputDir, "out.mp4"),
            fourcc=cv2.VideoWriter_fourcc(*"mp4v"),
            fps=fps,
            frameSize=size,
        )

        isFirstFrame = True
        thumbnail = None
        while cap.isOpened():
            res, frame = cap.read()
            if isFirstFrame:
                isFirstFrame = False
                thumbnail = frame

            if res == False:
                break

            resFram = inferenceImage(frame, threshold)
            result.write(resFram)

        cap.release()
        result.release()

        def createThumbnail(thumbnail):
            thumbnail = cv2.resize(
                src=thumbnail, dsize=(160, 160), interpolation=cv2.INTER_AREA
            )
            cv2.imwrite(os.path.join(inputDir, "thumbnail.jpg"), thumbnail)

        createThumbnail(thumbnail)

        async def uploadVideo():
            async with aiofiles.open(os.path.join(inputDir, "out.mp4"), "rb") as f:
                supabase.storage.from_("video").upload(
                    inputDir + ".mp4", await f.read(), {"content-type": "video/mp4"}
                )

        async def uploadThumbnail():
            async with aiofiles.open(
                os.path.join(inputDir, "thumbnail.jpg"), "rb"
            ) as f:
                supabase.storage.from_("thumbnail").upload(
                    inputDir + ".jpg", await f.read(), {"content-type": "image/jpeg"}
                )

        try:
            n = now()
            _, _ = await asyncio.gather(uploadVideo(), uploadThumbnail())
            print(now() - n)
        except Exception as e:
            print(e)

        updateArtifact(
            artifactId,
            {
                "status": "success",
                "path": "https://hdfxssmjuydwfwarxnfe.supabase.co/storage/v1/object/public/video/"
                + inputDir
                + ".mp4",
                "thumbnailURL": "https://hdfxssmjuydwfwarxnfe.supabase.co/storage/v1/object/public/thumbnail/"
                + inputDir
                + ".jpg",
            },
        )
    except:
        Process(
            updateArtifact(
                artifactId,
                {
                    "status": "fail",
                },
            )
        ).start()
    finally:
        shutil.rmtree(inputDir)


def updateArtifact(artifactId: str, body):
    artifact_ref = db.collection("artifacts").document(artifactId)
    artifact_snapshot = artifact_ref.get()
    if artifact_snapshot.exists:
        artifact_ref.update(body)
    # sendMessage(artifactId)


# This function cannot be automation test because the requirement of another device to receive notification
def sendMessage(artifactId: str, message: str = None):
    token = []
    artifact = db.collection("artifacts").document(artifactId).get()
    if not artifact.exists:
        return
    user_ref = db.collection("user").where(
        filter=FieldFilter("artifacts", "array-contains", "artifacts/" + artifactId)
    )
    for user in user_ref:
        token.append(user.get().to_dict()["deviceId"])
    if message is not None:
        messaging.MulticastMessage(
            data={
                "notification": {
                    "title": message,
                    "body": "Video "
                    + artifact.name
                    + " has done inference. Click here to see the video",
                },
            },
            android=messaging.AndroidConfig(
                notification=messaging.AndroidNotification(
                    icon="stock_ticker_update", color="#f45342"
                ),
            ),
        )
    else:
        messaging.MulticastMessage(
            data={
                "notification": {
                    "title": "Video " + artifact.name + " has done inference.",
                    "body": "Video "
                    + artifact.name
                    + " has done inference. Click here to see the video",
                },
            },
            android=messaging.AndroidConfig(
                notification=messaging.AndroidNotification(
                    icon="stock_ticker_update", color="#f45342"
                ),
            ),
        )
    response = messaging.send_multicast(message)
    return response.success_count