nicolasbuitragob commited on
Commit
fba6e1e
·
1 Parent(s): 0936f88
Files changed (4) hide show
  1. app.py +6 -7
  2. config.py +2 -1
  3. tasks.py +23 -15
  4. vitpose.py +2 -5
app.py CHANGED
@@ -1,11 +1,10 @@
1
  from fastapi import FastAPI, UploadFile, File, Response,Header, BackgroundTasks,Body
2
  from fastapi.staticfiles import StaticFiles
3
  from vitpose import VitPose
4
- import os
5
  from dotenv import load_dotenv
6
  from tasks import process_video
7
  from fastapi.responses import JSONResponse
8
- from config import API_KEY
9
  import logging
10
  logging.basicConfig(level=logging.INFO)
11
  logger = logging.getLogger(__name__)
@@ -13,7 +12,7 @@ app = FastAPI()
13
  vitpose = VitPose()
14
  # vitpose.pipeline.warmup()
15
 
16
- load_dotenv()
17
 
18
 
19
  app.mount("/static", StaticFiles())
@@ -30,10 +29,10 @@ def test():
30
  async def upload(background_tasks: BackgroundTasks,
31
  file: UploadFile = File(...),
32
  token: str = Header(...),
33
- user_id: str = Body(...)):
 
34
 
35
- if token != API_KEY:
36
- logger.info(f"Unauthorized {token} {API_KEY}")
37
  return JSONResponse(content={"message": "Unauthorized", "status": 401})
38
 
39
  logger.info("reading contents")
@@ -47,7 +46,7 @@ async def upload(background_tasks: BackgroundTasks,
47
  logger.info(f"file saved {file.filename}")
48
 
49
  # Create a clone of the file with content already read
50
- background_tasks.add_task(process_video, file.filename, vitpose, user_id)
51
 
52
  # Return the file as a response
53
  return JSONResponse(content={"message": "Video uploaded successfully", "status": 200})
 
1
  from fastapi import FastAPI, UploadFile, File, Response,Header, BackgroundTasks,Body
2
  from fastapi.staticfiles import StaticFiles
3
  from vitpose import VitPose
 
4
  from dotenv import load_dotenv
5
  from tasks import process_video
6
  from fastapi.responses import JSONResponse
7
+ from config import AI_API_TOKEN
8
  import logging
9
  logging.basicConfig(level=logging.INFO)
10
  logger = logging.getLogger(__name__)
 
12
  vitpose = VitPose()
13
  # vitpose.pipeline.warmup()
14
 
15
+ load_dotenv()
16
 
17
 
18
  app.mount("/static", StaticFiles())
 
29
  async def upload(background_tasks: BackgroundTasks,
30
  file: UploadFile = File(...),
31
  token: str = Header(...),
32
+ user_id: str = Body(...),
33
+ player_id: str = Body(...)):
34
 
35
+ if token != AI_API_TOKEN:
 
36
  return JSONResponse(content={"message": "Unauthorized", "status": 401})
37
 
38
  logger.info("reading contents")
 
46
  logger.info(f"file saved {file.filename}")
47
 
48
  # Create a clone of the file with content already read
49
+ background_tasks.add_task(process_video, file.filename, vitpose, user_id,player_id)
50
 
51
  # Return the file as a response
52
  return JSONResponse(content={"message": "Video uploaded successfully", "status": 200})
config.py CHANGED
@@ -4,4 +4,5 @@ from dotenv import load_dotenv
4
  load_dotenv()
5
 
6
  API_URL = os.getenv("API_URL")
7
- API_KEY = os.getenv("API_KEY")
 
 
4
  load_dotenv()
5
 
6
  API_URL = os.getenv("API_URL")
7
+ API_KEY = os.getenv("API_KEY")
8
+ AI_API_TOKEN = os.getenv("AI_API_TOKEN")
tasks.py CHANGED
@@ -1,38 +1,46 @@
1
  from vitpose import VitPose
2
  import requests
3
  import os
 
4
  from fastapi import UploadFile
5
- from config import API_URL
6
- import time
7
  import logging
8
  logging.basicConfig(level=logging.INFO)
9
  logger = logging.getLogger(__name__)
10
 
11
- def process_video(video_path: str,vitpose: VitPose,user_id: str):
12
 
13
- logger.info(f"starting task {video_path}")
 
 
 
 
 
14
 
 
 
 
15
 
16
- new_file_name = video_path.split(".")[0] + '.' + video_path.split(".")[1]
17
- new_file_name = os.path.join("static", new_file_name)
18
  logger.info(f"new file name {new_file_name}")
19
 
20
  vitpose.output_video_path = new_file_name
21
  annotated_frames = vitpose.run(video_path)
22
 
 
23
 
 
24
 
25
-
26
- annotated_video_path = vitpose.frames_to_video(annotated_frames)
27
-
28
- logger.info(f"Video processed {annotated_video_path}")
29
-
30
- with open(annotated_video_path, "rb") as f:
31
  contents = f.read()
32
 
33
  url = API_URL+ "/excercises/webhooks/video-processed"
34
  logger.info(f"Sending video to {url}")
35
- files = {"file": (annotated_video_path, contents, "video/mp4")}
36
- response = requests.post(url, files=files, data={"user_id":user_id,"typeMessage":"video_processed","file_name":annotated_video_path}, stream=True)
 
 
 
 
 
37
  logger.info(f"Video sent to {url}")
38
- print(response.json())
 
1
  from vitpose import VitPose
2
  import requests
3
  import os
4
+ from config import API_URL,API_KEY
5
  from fastapi import UploadFile
 
 
6
  import logging
7
  logging.basicConfig(level=logging.INFO)
8
  logger = logging.getLogger(__name__)
9
 
10
+ def process_video(file_name: str,vitpose: VitPose,user_id: str,player_id: str):
11
 
12
+ video_path = file_name
13
+
14
+ contents = open(video_path, "rb").read()
15
+
16
+ with open(video_path, "wb") as f:
17
+ f.write(contents)
18
 
19
+ logger.info(f"file saved {video_path}")
20
+
21
+ logger.info(f"starting task {video_path}")
22
 
23
+ new_file_name = os.path.join("static", video_path)
 
24
  logger.info(f"new file name {new_file_name}")
25
 
26
  vitpose.output_video_path = new_file_name
27
  annotated_frames = vitpose.run(video_path)
28
 
29
+ vitpose.frames_to_video(annotated_frames)
30
 
31
+ logger.info(f"Video processed {video_path}")
32
 
33
+ with open(new_file_name, "rb") as f:
 
 
 
 
 
34
  contents = f.read()
35
 
36
  url = API_URL+ "/excercises/webhooks/video-processed"
37
  logger.info(f"Sending video to {url}")
38
+ files = {"file": (video_path, contents, "video/mp4")}
39
+
40
+ response = requests.post(url, files=files,
41
+ data={"user_id":user_id,"typeMessage":"video_processed","file_name":video_path,
42
+ "player_id":player_id},
43
+ stream=True,
44
+ headers={"token":API_KEY})
45
  logger.info(f"Video sent to {url}")
46
+
vitpose.py CHANGED
@@ -96,7 +96,6 @@ class VitPose:
96
  else:
97
  print(f"Dimensions: {width}x{height}")
98
  out = cv2.VideoWriter(self.output_video_path, fourcc, self.video_metadata["fps"], (width, height))
99
-
100
  for frame in frames:
101
  if rotate:
102
  # Rotate landscape videos 90 degrees to make them vertical
@@ -104,7 +103,5 @@ class VitPose:
104
  out.write(rotated_frame)
105
  else:
106
  # Already vertical, no rotation needed
107
- out.write(frame)
108
-
109
- out.release()
110
- return self.output_video_path
 
96
  else:
97
  print(f"Dimensions: {width}x{height}")
98
  out = cv2.VideoWriter(self.output_video_path, fourcc, self.video_metadata["fps"], (width, height))
 
99
  for frame in frames:
100
  if rotate:
101
  # Rotate landscape videos 90 degrees to make them vertical
 
103
  out.write(rotated_frame)
104
  else:
105
  # Already vertical, no rotation needed
106
+ out.write(frame)
107
+ out.release()