Update app.py

app.py

@@ -6,7 +6,9 @@ from pathlib import Path
 import shutil
 import logging
 import uvicorn
-from typing import Optional
+from typing import Optional, List
+import pandas as pd
+import json
 
 from configs import ModelConfig, InferenceConfig
 from tools.models import load_pipeline
@@ -19,12 +21,29 @@ app = FastAPI(title="Sign Language Recognition API")
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Define a Pydantic model for individual prediction
+class Prediction(BaseModel):
+    gloss: str
+    score: float
+    start_time: float
+    end_time: float
+    inference_time: float
+
 # Define a Pydantic model for the response
 class InferenceResponse(BaseModel):
     status: str
-    predictions: Optional[list] = None
+    predictions: Optional[List[Prediction]] = None
     message: Optional[str] = None
 
+# Define id2gloss mapping
+# Đây là một ví dụ. Bạn cần thay thế bằng bản đồ thực tế của bạn.
+id2gloss = {
+    "0": "hello",
+    "1": "thanks",
+    "2": "yes",
+    # Thêm các ánh xạ cần thiết
+}
+
 @app.post("/inference", response_model=InferenceResponse)
 async def inference_endpoint(
     file: UploadFile = File(...),
@@ -42,61 +61,106 @@ async def inference_endpoint(
     Returns:
         InferenceResponse: Kết quả nhận diện.
     """
-    # Kiểm tra file có hợp lệ không
-    if not file.filename.endswith((".mp4", ".avi", ".mov", ".mkv")):
-        raise HTTPException(status_code=400, detail="Unsupported file type.")
-    
-    # Tạo thư mục output nếu không tồn tại
-    output_path = Path(output_dir)
-    output_path.mkdir(parents=True, exist_ok=True)
-    
-    # Lưu video tạm thời
-    video_path = output_path / file.filename
-    with open(video_path, "wb") as buffer:
-        shutil.copyfileobj(file.file, buffer)
-    
-    logger.info(f"Video saved to {video_path}")
-
-    # Tải cấu hình mô hình dựa trên model_name
     try:
-        if model_name == "spoter":
-            model_config = ModelConfig(arch="spoter", pretrained="vsltranslation/spoter_v3.0")
-        elif model_name == "sl_gcn":
-            model_config = ModelConfig(arch="sl_gcn", pretrained="vsltranslation/sl_gcn_joint_v3_0")
-        elif model_name == "dsta_slr":
-            model_config = ModelConfig(arch="dsta_slr", pretrained="models/dsta_slr_joint_motion_v3_0.onnx")
-        else:
-            raise ValueError("Unsupported model name.")
+        # Kiểm tra file có hợp lệ không
+        if not file.filename.endswith((".mp4", ".avi", ".mov", ".mkv")):
+            raise HTTPException(status_code=400, detail="Unsupported file type.")
         
-        inference_config = InferenceConfig(
-            source=str(video_path),
-            output_dir=str(output_path),
-            use_onnx=True if model_config.pretrained.endswith(".onnx") else False,
-            device="cpu",  # Bạn có thể thay đổi thành "cuda" nếu sử dụng GPU
-            cache_dir="models/huggingface",
-            visualize=False,
-            show_skeleton=False,
-            visibility=0.5,
-            angle_threshold=140,
-            min_num_up_frames=10,
-            min_num_down_frames=10,
-            delay=400,
-            top_k=3,
-            bone_stream=False,
-            motion_stream=False
-        )
+        # Tạo thư mục output nếu không tồn tại
+        output_path = Path(output_dir)
+        output_path.mkdir(parents=True, exist_ok=True)
         
-        # Tải pipeline
-        pipeline = load_pipeline(model_config, inference_config)
-        logger.info("Pipeline loaded successfully.")
+        # Lưu video tạm thời
+        video_path = output_path / file.filename
+        with open(video_path, "wb") as buffer:
+            shutil.copyfileobj(file.file, buffer)
+        
+        logger.info(f"Video saved to {video_path}")
+
+        # Tải cấu hình mô hình dựa trên model_name
+        try:
+            if model_name == "spoter":
+                model_config = ModelConfig(arch="spoter", pretrained="vsltranslation/spoter_v3.0")
+            elif model_name == "sl_gcn":
+                model_config = ModelConfig(arch="sl_gcn", pretrained="vsltranslation/sl_gcn_joint_v3_0")
+            elif model_name == "dsta_slr":
+                model_config = ModelConfig(arch="dsta_slr", pretrained="models/dsta_slr_joint_motion_v3_0.onnx")
+            else:
+                raise ValueError("Unsupported model name.")
+            
+            inference_config = InferenceConfig(
+                source=str(video_path),
+                output_dir=str(output_path),
+                use_onnx=True if model_config.pretrained.endswith(".onnx") else False,
+                device="cpu",  # Bạn có thể thay đổi thành "cuda" nếu sử dụng GPU
+                cache_dir="models/huggingface",
+                visualize=False,
+                show_skeleton=False,
+                visibility=0.5,
+                angle_threshold=140,
+                min_num_up_frames=10,
+                min_num_down_frames=10,
+                delay=400,
+                top_k=3,
+                bone_stream=False,
+                motion_stream=True  # Theo cấu hình YAML bạn cung cấp
+            )
+            
+            # Tải pipeline hoặc session
+            pipeline_or_session = load_pipeline(model_config, inference_config)
+            logger.info("Pipeline loaded successfully.")
+
+            # Chạy inference
+            run_inference(model_config, inference_config, pipeline_or_session)
+            logger.info("Inference completed successfully.")
 
-        # Chạy inference
-        result = run_inference(model_config, inference_config, pipeline)
-        logger.info("Inference completed successfully.")
+            # Đọc kết quả từ results.csv
+            results_csv = output_path / "results.csv"
+            if not results_csv.exists():
+                raise HTTPException(status_code=500, detail="Inference did not produce results.")
 
-        # Trả về kết quả
-        return InferenceResponse(status="success", predictions=result.get("results", []))
+            results_df = pd.read_csv(results_csv)
 
+            # Chuyển đổi DataFrame thành list of Prediction
+            predictions = []
+            for _, row in results_df.iterrows():
+                # Giả sử results.csv có các cột: start_time, end_time, inference_time, prediction
+                # Và 'prediction' là một danh sách các từ điển với 'gloss' và 'score'
+                start_time = row.get("start_time", 0.0)
+                end_time = row.get("end_time", 0.0)
+                inference_time = row.get("inference_time", 0.0)
+                prediction_list = row.get("prediction", [])
+                
+                if isinstance(prediction_list, str):
+                    # Nếu prediction được lưu dưới dạng chuỗi JSON
+                    try:
+                        prediction_list = json.loads(prediction_list.replace("'", '"'))
+                    except json.JSONDecodeError:
+                        prediction_list = []
+                
+                for pred in prediction_list:
+                    gloss = pred.get('gloss', 'Unknown')
+                    score = pred.get('score', 0.0)
+                    predictions.append(Prediction(
+                        gloss=gloss,
+                        score=score,
+                        start_time=start_time,
+                        end_time=end_time,
+                        inference_time=inference_time
+                    ))
+            
+            return InferenceResponse(status="success", predictions=predictions)
+
+        except ValueError as ve:
+            logger.exception("ValueError during inference.")
+            raise HTTPException(status_code=400, detail=str(ve))
+        except Exception as e:
+            logger.exception("Error during inference.")
+            raise HTTPException(status_code=500, detail=str(e))
+    
     except Exception as e:
-        logger.exception("Error during inference.")
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.exception("Unexpected error.")
+        raise HTTPException(status_code=500, detail="An unexpected error occurred.")
+
+if __name__ == "__main__":
+    uvicorn.run("api.main:app", host="0.0.0.0", port=8000, reload=True)