webhooks and background task

app.py

@@ -1,6 +1,19 @@
-from fastapi import FastAPI
-
+from fastapi import FastAPI, UploadFile, File, Response,Header, BackgroundTasks,Body
+from fastapi.staticfiles import StaticFiles
+from vitpose import VitPose
+import os 
+from dotenv import load_dotenv
+from tasks import process_video
+from fastapi.responses import JSONResponse
+from config import API_KEY
 app = FastAPI()
+vitpose = VitPose()
+# vitpose.pipeline.warmup()
+
+load_dotenv()
+
+
+app.mount("/static", StaticFiles())
 
 @app.get("/")
 def read_root():
@@ -8,4 +21,27 @@ def read_root():
 
 @app.get("/test")
 def test():
-    return {"message": "from test"}
+    return {"message": "from test"}
+
+@app.post("/upload")
+async def upload(background_tasks: BackgroundTasks,
+                 file: UploadFile = File(...), 
+                 token: str = Header(...),
+                 user_id: str = Body(...)):
+    
+    if token != API_KEY:
+        return JSONResponse(content={"message": "Unauthorized", "status": 401})
+    
+    contents = await file.read()
+    # Save the file to the local directory  
+    with open(file.filename, "wb") as f:
+        f.write(contents)
+    
+
+    
+    # Create a clone of the file with content already read
+    background_tasks.add_task(process_video, file.filename, vitpose, user_id)  
+
+    # Return the file as a response
+    return JSONResponse(content={"message": "Video uploaded successfully", "status": 200})
+    

config.py

@@ -0,0 +1,7 @@
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+API_URL = os.getenv("API_URL")
+API_KEY = os.getenv("API_KEY")

tasks.py

@@ -0,0 +1,26 @@
+from vitpose import VitPose
+import requests
+import os
+from fastapi import UploadFile
+from config import API_URL
+import time
+
+def process_video(video_path: str,vitpose: VitPose,user_id: str):
+        
+    new_file_name = video_path.split(".")[0] + "edited." + video_path.split(".")[1]
+    new_file_name = os.path.join("static", new_file_name)
+    
+    vitpose.output_video_path = new_file_name
+    annotated_frames = vitpose.run(video_path)
+    annotated_video_path = vitpose.frames_to_video(annotated_frames,rotate=True)
+    
+    with open(annotated_video_path, "rb") as f:
+        contents = f.read()
+        
+    url = API_URL+ "/excercises/webhooks/video-processed"
+    
+    files = {"file": (annotated_video_path, contents, "video/mp4")}
+    response = requests.post(url, files=files, data={"user_id":user_id,"typeMessage":"video_processed","file_name":annotated_video_path}, stream=True)
+    print(response.json())
+    os.remove(video_path)
+    os.remove(annotated_video_path)

vitpose.py

@@ -0,0 +1,101 @@
+import torch
+from rt_pose import PoseEstimationPipeline
+import cv2
+import supervision as sv
+import numpy as np
+from rt_pose import PoseEstimationPipeline, PoseEstimationOutput
+
+
+class VitPose:
+    def __init__(self):
+        self.pipeline = PoseEstimationPipeline(
+            object_detection_checkpoint="PekingU/rtdetr_r50vd_coco_o365",
+            pose_estimation_checkpoint="usyd-community/vitpose-plus-small",
+            device="cuda" if torch.cuda.is_available() else "cpu",
+            dtype=torch.bfloat16,
+            compile=True,  # or True to get more speedup
+        )
+        self.output_video_path = None
+        self.video_metadata = {}
+        
+        
+    def video_to_frames(self,video):
+        frames = []
+        cap = cv2.VideoCapture(video)
+        self.video_metadata = {
+            "fps": cap.get(cv2.CAP_PROP_FPS),
+            "width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
+            "height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
+        }
+        
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frames.append(frame)
+        return frames[:10]
+    
+    def run(self,video):
+        frames = self.video_to_frames(video)
+        annotated_frames = []
+        for frame in frames:
+            output = self.pipeline(frame)
+            annotated_frame = self.visualize_output(frame,output)
+            annotated_frames.append(annotated_frame)
+        return annotated_frames
+    
+    
+ 
+    def visualize_output(self,image: np.ndarray, output: PoseEstimationOutput, confidence: float = 0.3) -> np.ndarray:
+        """
+        Visualize pose estimation output.
+        """
+        keypoints_xy = output.keypoints_xy.float().cpu().numpy()
+        scores = output.scores.float().cpu().numpy()
+
+        # Supervision will not draw vertices with `0` score
+        # and coordinates with `(0, 0)` value
+        invisible_keypoints = scores < confidence
+        scores[invisible_keypoints] = 0
+        keypoints_xy[invisible_keypoints] = 0
+
+        keypoints = sv.KeyPoints(xy=keypoints_xy, confidence=scores)
+
+        _, y_min, _, y_max = output.person_boxes_xyxy.T
+        height = int((y_max - y_min).mean().item())
+        radius = max(height // 100, 4)
+        thickness = max(height // 200, 2)
+        edge_annotator = sv.EdgeAnnotator(color=sv.Color.YELLOW, thickness=thickness)
+        vertex_annotator = sv.VertexAnnotator(color=sv.Color.ROBOFLOW, radius=radius)
+
+        annotated_frame = image.copy()
+        annotated_frame = edge_annotator.annotate(annotated_frame, keypoints)
+        annotated_frame = vertex_annotator.annotate(annotated_frame, keypoints)
+
+        return annotated_frame
+    
+    def frames_to_video(self, frames, rotate=False):
+        
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        height = self.video_metadata["height"]
+        width = self.video_metadata["width"]
+        
+        # If rotation is requested, swap dimensions for the output video
+        if rotate:
+            print(f"Original dimensions: {width}x{height}, Rotated dimensions: {height}x{width}")
+            # For the VideoWriter, we need to specify the dimensions of the output frames
+            out = cv2.VideoWriter(self.output_video_path, fourcc, self.video_metadata["fps"], (height, width))
+        else:
+            print(f"Dimensions: {width}x{height}")
+            out = cv2.VideoWriter(self.output_video_path, fourcc, self.video_metadata["fps"], (width, height))
+        
+        for frame in frames:
+            if rotate:
+                # Rotate 90 degrees clockwise
+                rotated_frame = cv2.rotate(frame, cv2.ROTATE_90_COUNTERCLOCKWISE)
+                out.write(rotated_frame)
+            else:
+                out.write(frame)
+            
+        out.release()
+        return self.output_video_path