app.py

import logging
from time import time
import pandas as pd
import numpy as np
import cv2
from typing import Optional
from pathlib import Path
from fastapi import FastAPI, HTTPException, UploadFile, File, Query
from fastapi.responses import JSONResponse
import mediapipe as mp

from configs import ModelConfig, InferenceConfig
from utils import config_logger, POSE_BASED_MODELS
from data import Arm, get_sample_timestamp, ok_to_get_frame
from tools import load_pipeline, Predictions
from visualization import draw_text_on_image

app = FastAPI()

# Định nghĩa ba preset model
MODEL_PRESETS = {
    "dsta_slr": {
        "model": ModelConfig(
            arch="dsta_slr",
            pretrained="models/dsta_slr_joint_motion_v3_0.onnx",
        ),
        "inference": InferenceConfig(
            source="upload",  # Sử dụng upload, không webcam
            output_dir="demo/run_1",
            use_onnx=True,
            show_skeleton=True,
            visualize=True,
            bone_stream=False,
            motion_stream=True,
        ),
    },
    "sl_gcn": {
        "model": ModelConfig(
            arch="sl_gcn",
            pretrained="models/dsta_slr_joint_motion_v3_0.onnx",
        ),
        "inference": InferenceConfig(
            source="upload",
            output_dir="demo/run_1",
            use_onnx=True,
            show_skeleton=True,
            visualize=True,
            bone_stream=True,
            motion_stream=False,
        ),
    },
    "spoter": {
        "model": ModelConfig(
            arch="spoter",
            pretrained="models/spoter_v3.0.onnx",
        ),
        "inference": InferenceConfig(
            source="upload",
            output_dir="demo/run_1",
            use_onnx=True,
            show_skeleton=True,
            visualize=True,
        ),
    },
}

config_logger("inference.log")
logging.info("API started")

SPOTER_POSE_LANDMARKS = [
    mp.solutions.pose.PoseLandmark.NOSE,
    mp.solutions.pose.PoseLandmark.LEFT_EYE, 
    mp.solutions.pose.PoseLandmark.RIGHT_EYE, 
    mp.solutions.pose.PoseLandmark.RIGHT_SHOULDER,
    mp.solutions.pose.PoseLandmark.LEFT_SHOULDER,
    mp.solutions.pose.PoseLandmark.RIGHT_ELBOW,
    mp.solutions.pose.PoseLandmark.LEFT_ELBOW,
    mp.solutions.pose.PoseLandmark.RIGHT_WRIST,
    mp.solutions.pose.PoseLandmark.LEFT_WRIST 
]

SPOTER_HAND_LANDMARKS = [
    mp.solutions.hands.HandLandmark.WRIST,
    mp.solutions.hands.HandLandmark.INDEX_FINGER_TIP, mp.solutions.hands.HandLandmark.INDEX_FINGER_DIP, 
    mp.solutions.hands.HandLandmark.INDEX_FINGER_PIP, mp.solutions.hands.HandLandmark.INDEX_FINGER_MCP,
    mp.solutions.hands.HandLandmark.MIDDLE_FINGER_TIP, mp.solutions.hands.HandLandmark.MIDDLE_FINGER_DIP, 
    mp.solutions.hands.HandLandmark.MIDDLE_FINGER_PIP, mp.solutions.hands.HandLandmark.MIDDLE_FINGER_MCP,
    mp.solutions.hands.HandLandmark.RING_FINGER_TIP, mp.solutions.hands.HandLandmark.RING_FINGER_DIP, 
    mp.solutions.hands.HandLandmark.RING_FINGER_PIP, mp.solutions.hands.HandLandmark.RING_FINGER_MCP,
    mp.solutions.hands.HandLandmark.PINKY_TIP, mp.solutions.hands.HandLandmark.PINKY_DIP, 
    mp.solutions.hands.HandLandmark.PINKY_PIP, mp.solutions.hands.HandLandmark.PINKY_MCP,
    mp.solutions.hands.HandLandmark.THUMB_TIP, mp.solutions.hands.HandLandmark.THUMB_IP, 
    mp.solutions.hands.HandLandmark.THUMB_MCP, mp.solutions.hands.HandLandmark.THUMB_CMC,
]


@app.get("/healthcheck")
async def healthcheck():
    return JSONResponse(status_code=200, content={"status": "UP"})


def run_inference(model_config, inference_config, input_frames):
    pipeline = load_pipeline(model_config, inference_config)
    logging.info("Pipeline loaded")

    right_arm = Arm("right", inference_config.visibility)
    left_arm = Arm("left", inference_config.visibility)
    data = []
    results = None
    predictions = Predictions()

    mp_holistic = mp.solutions.holistic
    mp_drawing = mp.solutions.drawing_utils
    mp_drawing_styles = mp.solutions.drawing_styles

    custom_pose_style = mp_drawing_styles.get_default_pose_landmarks_style()
    custom_right_hand_style = mp_drawing_styles.get_default_hand_landmarks_style()
    custom_left_hand_style = mp_drawing_styles.get_default_hand_landmarks_style()
    custom_pose_connections = list(mp_holistic.POSE_CONNECTIONS)
    custom_hand_connections = list(mp_holistic.HAND_CONNECTIONS)

    if inference_config.show_skeleton:
        pose_landmarks = SPOTER_POSE_LANDMARKS
        hand_landmarks = SPOTER_HAND_LANDMARKS
        for landmark in mp.solutions.pose.PoseLandmark:
            if landmark in pose_landmarks:
                custom_pose_style[landmark] = mp.drawing.DrawingSpec(color=(0,255,0), thickness=2, circle_radius=2)
            else:
                custom_pose_style[landmark] = mp.drawing.DrawingSpec(color=(0,0,0), thickness=0, circle_radius=0) 
                for connection_tuple in custom_pose_connections:
                    if landmark.value in connection_tuple:
                        custom_pose_connections.remove(connection_tuple)
        for landmark in mp.solutions.hands.HandLandmark:
            if landmark in hand_landmarks:
                custom_right_hand_style[landmark] = mp.drawing.DrawingSpec(color=(0,0,255), thickness=2, circle_radius=2)
                custom_left_hand_style[landmark] = mp.drawing.DrawingSpec(color=(255,0,0), thickness=2, circle_radius=2)
            else:
                custom_right_hand_style[landmark] = mp.drawing.DrawingSpec(color=(0,0,0), thickness=0, circle_radius=0)
                custom_left_hand_style[landmark] = mp.drawing.DrawingSpec(color=(0,0,0), thickness=0, circle_radius=0)
                for connection_tuple in custom_hand_connections:
                    if landmark.value in connection_tuple:
                        custom_hand_connections.remove(connection_tuple)

    writer = None
    if inference_config.output_dir is not None:
        out_path = Path(inference_config.output_dir)
        out_path.mkdir(parents=True, exist_ok=True)
        if len(input_frames) > 0 and isinstance(input_frames[0], np.ndarray):
            h, w, _ = input_frames[0].shape
            writer = cv2.VideoWriter(str(out_path / "output.mp4"), cv2.VideoWriter_fourcc(*"mp4v"), 30, (w, h))

    with mp_holistic.Holistic(min_detection_confidence=0.9, min_tracking_confidence=0.5) as holistic:
        # giả định mỗi frame ~33ms, ở đây chỉ là demo logic
        current_time_ms = 0
        for frame in input_frames:
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            rgb_frame.flags.writeable = False
            detection_results = holistic.process(rgb_frame)

            try:
                landmarks = detection_results.pose_landmarks.landmark
            except:
                current_time_ms += 33
                continue

            left_arm.set_pose(landmarks)
            right_arm.set_pose(landmarks)

            left_arm_ok_to_get_frame = ok_to_get_frame(
                arm=left_arm,
                angle_threshold=inference_config.angle_threshold,
                min_num_up_frames=inference_config.min_num_up_frames,
                min_num_down_frames=inference_config.min_num_down_frames,
                current_time=current_time_ms,
                delay=inference_config.delay,
            )
            right_arm_ok_to_get_frame = ok_to_get_frame(
                arm=right_arm,
                angle_threshold=inference_config.angle_threshold,
                min_num_up_frames=inference_config.min_num_up_frames,
                min_num_down_frames=inference_config.min_num_down_frames,
                current_time=current_time_ms,
                delay=inference_config.delay,
            )

            if left_arm_ok_to_get_frame or right_arm_ok_to_get_frame:
                predictions = Predictions()
                data.append(detection_results if inference_config.use_pose_model else frame)

            start_time, end_time = get_sample_timestamp(left_arm, right_arm)
            start_time /= 1000
            end_time /= 1000

            if start_time != 0 and end_time != 0:
                start_inference_time = time()
                predictions = Predictions(predictions=pipeline(np.array(data)))
                predictions.inference_time = time() - start_inference_time
                predictions.start_time = start_time
                predictions.end_time = end_time
                logging.info(str(predictions))
                results = predictions.merge_results(results)

                # Reset
                start_time = 0
                end_time = 0
                left_arm.reset_state()
                right_arm.reset_state()
                data = []

            # Vẽ kết quả
            frame = left_arm.visualize(frame, (20, 10), "Left arm angle")
            frame = right_arm.visualize(frame, (20, 40), "Right arm angle")
            frame = predictions.visualize(frame, (20, 70))

            if inference_config.show_skeleton:
                mp.drawing.draw_landmarks(
                    frame,
                    detection_results.pose_landmarks,
                    connections=custom_pose_connections,
                    landmark_drawing_spec=custom_pose_style
                )
                mp.drawing.draw_landmarks(
                    frame,
                    detection_results.right_hand_landmarks,
                    connections=custom_hand_connections,
                    landmark_drawing_spec=custom_right_hand_style
                )
                mp.drawing.draw_landmarks(
                    frame,
                    detection_results.left_hand_landmarks,
                    connections=custom_hand_connections,
                    landmark_drawing_spec=custom_left_hand_style
                )

            if writer is not None:
                writer.write(frame)

            current_time_ms += 33

    if writer is not None:
        writer.release()
    if results is not None:
        pd.DataFrame(results).to_csv(Path(inference_config.output_dir) / "results.csv", index=False)

    return predictions.predictions, results


@app.post("/inference")
async def inference_endpoint(
    model_name: str = Query(..., description="Choose model: dsta_slr, sl_gcn, spoter"),
    output_option: str = Query("all", description="Output option: 'predictions', 'csv', 'video', 'all'"),
    output_dir: str = Query("demo/run_1", description="Output directory for results"),
    file: UploadFile = File(...)
):
    """
    Inference endpoint:
    - model_name: chọn mô hình: dsta_slr, sl_gcn, spoter
    - output_option: 'predictions', 'csv', 'video', hoặc 'all'
    - output_dir: thư mục output, vd: 'my_results'
    - file: upload 1 file video
    """

    if model_name not in MODEL_PRESETS:
        raise HTTPException(status_code=400, detail="Invalid model_name")

    # Đọc video từ file upload
    video_bytes = np.asarray(bytearray(await file.read()), dtype=np.uint8)
    temp_video_path = Path("temp_input.mp4")
    with open(temp_video_path, "wb") as f:
        f.write(video_bytes)
    cap = cv2.VideoCapture(str(temp_video_path))

    input_frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        input_frames.append(frame)
    cap.release()

    # Load config từ preset
    model_config = MODEL_PRESETS[model_name]["model"]
    inference_config = MODEL_PRESETS[model_name]["inference"]

    # Ghi đè output_dir theo yêu cầu người dùng
    inference_config.output_dir = output_dir

    if model_config.arch in POSE_BASED_MODELS:
        inference_config.use_pose_model = True
    else:
        inference_config.use_pose_model = False

    predictions, results = run_inference(model_config, inference_config, input_frames)

    resp = {}
    out_dir = Path(inference_config.output_dir)
    if predictions is None:
        predictions = []

    if output_option in ["predictions", "all"]:
        resp["predictions"] = predictions

    if output_option in ["csv", "all"]:
        csv_path = str(out_dir / "results.csv")
        resp["csv_path"] = csv_path if Path(csv_path).exists() else None

    if output_option in ["video", "all"]:
        video_path = str(out_dir / "output.mp4")
        resp["video_path"] = video_path if Path(video_path).exists() else None

    return resp