Browse Source

almost work

LeoDu 4 years ago
parent
commit
b2af0c7427

+ 1 - 1
backend_refactor/config/__init__.py

@@ -1,7 +1,7 @@
 import os
 
 VERSION_PREFIX = "/v1/"
-UPLOAD_DIR = "upload/"
+UPLOAD_DIR = "/data/"
 MAX_CONTENT_LENGTH = 100 * 1024 * 1024
 
 MYSQL_URI = f'mysql://{os.environ["MYSQL_USER"]}:{os.environ["MYSQL_PASSWORD"]}@{os.environ["MYSQL_ADDRESS"]}/{os.environ["MYSQL_DATABASE"]}'

+ 1 - 1
backend_refactor/db/task.py

@@ -3,7 +3,7 @@ import time
 
 class Task(db.Model):
     uuid = db.Column(db.String(64), primary_key=True)
-    type = db.Column(db.Enum('PHOTO_STANDINGPOSE', 'VIDEO_HIGHKNEES', 'VIDEO_CHRUNCH', 'VIDEO_STANDINGLONGJUMP', 'VIDEO_PLANK', 'VIDEO_SQUAT'))
+    type = db.Column(db.Enum('PHOTO_STANDINGPOSE', 'VIDEO_HIGHKNEES', 'VIDEO_CRUNCH', 'VIDEO_STANDINGLONGJUMP', 'VIDEO_PLANK', 'VIDEO_SQUAT'))
     # 站姿 高抬腿 卷腹 立定跳远 平板支撑 深蹲
     create_time = db.Column(db.BigInteger)
     status = db.Column(db.Enum('QUEUEING','RUNNING','FINISHED','ERROR'))

+ 21 - 2
backend_refactor/route/analyze.py

@@ -2,8 +2,12 @@ from hashlib import new
 import random
 import logging
 import uuid, json
+from db import task
+from service.photo.BasePhotoAnalyzer import BasePhotoAnalyzer
+from service.video.BaseVideoAnalyzer import BaseVideoAnalyzer
 from config import MYSQL_URI
 # from service.video.Video3DAnalyzer import Video3DAnalyzer
+
 from db import db
 from flask import request
 from db.task import Task
@@ -14,6 +18,12 @@ from service.photo.StandingPhotoAnalyzer import StandingPhotoAnalyzer
 from sqlalchemy import create_engine
 
 
+import logging
+
+logging.basicConfig(format='%(asctime)s - %(name)s[line:%(lineno)d] - %(levelname)s: %(message)s',
+                    level=logging.DEBUG)
+
+
 bp = Blueprint("analyze", __name__, url_prefix='/')
 
 def callback_generator(task_uuid):
@@ -29,7 +39,7 @@ def callback_generator(task_uuid):
                     logging.error("error analyzing: ", exc_info=kwargs['error'])
                 
                 db.session.commit()
-                logging.trace("calling to" + status)
+                logging.debug("calling to" + status)
         except Exception as e:
             logging.debug("nmsl", exc_info=e)
     return callback
@@ -58,8 +68,17 @@ def video_analysis():
     req = request.get_json()
     
     # ana_3d = Video3DAnalyzer()
+    ana = BaseVideoAnalyzer(video_uuid=req['video_uuid'])
 
     task_uuid = str(uuid.uuid4())
-    new_task = Task(task_uuid, req['type'])
+    new_task = Task(task_uuid, "VIDEO_CRUNCH")
     db.session.add(new_task)
     db.session.commit()
+
+    pool.submit(ana.analyze, callback_generator(task_uuid))
+
+    return make_response(201, "success", {'task_uuid' : task_uuid})
+
+
+
+

+ 2 - 0
backend_refactor/service/aiutil/torch_openpose.py

@@ -6,6 +6,7 @@ import torch
 
 from service.aiutil import util
 from service.aiutil.model import bodypose_model,bodypose_25_model
+import logging
 
 model_coco = 'model/body_coco.pth'
 model_body25 = 'model/body_25.pth'
@@ -23,6 +24,7 @@ class torch_openpose(object):
             self.njoint = 19
             self.npaf = 38
             self.model.load_state_dict(torch.load(model_coco))
+        logging.debug("CUDA: " + str(torch.cuda.is_available()))
         if torch.cuda.is_available():
             self.model = self.model.cuda()
         self.model.eval()

+ 75 - 4
backend_refactor/service/video/BaseVideoAnalyzer.py

@@ -1,6 +1,77 @@
-class BaseVideoAnalyzer:
-    def __init__(self):
-        pass
 
-    def analyze(self, filename):
+import os
+from typing import Callable
+from config import UPLOAD_DIR
+import cv2
+import logging
+import numpy as np
+
+
+def split_video(video_uuid):
+    video_path = UPLOAD_DIR + video_uuid + ".mp4"
+    frames_dir = UPLOAD_DIR + "frames/" + video_uuid + "/"
+
+    if os.path.exists(frames_dir):  # Already split
         return
+
+    os.makedirs(frames_dir)
+
+    cap = cv2.VideoCapture(video_path)
+    FPS = cap.get(5)  # 5 means fps
+    logging.debug("Splitting video %s, FPS: %s" % (video_path, FPS))
+
+    c = 1
+    sampling_fps = 1
+
+    while True:
+        ret, frame = cap.read()
+        if ret:
+            if c % sampling_fps == 0:
+                target_path = frames_dir + str(c) + ".jpg"
+                cv2.imwrite(target_path, frame)
+                logging.debug("Writing frame %s, %s" % (c, target_path))
+            c += 1
+        else:
+            break
+
+    cap.release()
+
+
+from service.aiutil.torch_openpose import torch_openpose
+from service.aiutil.util import draw_bodypose
+
+
+def run_openpose_for_frames(video_uuid):
+    tp = torch_openpose('body_25')
+
+    frames_dir = UPLOAD_DIR + "frames/" + video_uuid + "/"
+    for i in os.listdir(frames_dir):
+        if i.endswith(".jpg") and "result" not in i:
+            
+            logging.debug("Running openpose for %s %s" % (video_uuid, i))
+
+            oriImg = cv2.imread(frames_dir + i)
+            poses = tp(oriImg)
+            np.save(frames_dir + i + ".npy", poses)
+
+            canvas = draw_bodypose(oriImg, poses, 'body_25')
+
+            cv2.imwrite(frames_dir + i.replace(".jpg", ".result.jpg"), canvas)
+
+
+
+class BaseVideoAnalyzer:
+    def __init__(self, video_uuid):
+        self.video_uuid = video_uuid
+
+    def analyze(self, callback: Callable):
+        try:
+            callback("RUNNING")
+            split_video(self.video_uuid)
+            run_openpose_for_frames(self.video_uuid)
+            callback("FINISHED", result={})
+        except Exception as e:
+            callback("ERROR", error=e)
+            raise e
+
+

+ 0 - 50
backend_refactor/service/video/Video3DAnalyzer.py

@@ -1,6 +1,5 @@
 import glob
 import logging
-import os
 import subprocess as sp
 import time
 
@@ -11,7 +10,6 @@ import numpy as np
 # from detectron2.engine import DefaultPredictor
 # from detectron2.utils.logger import setup_logger
 
-from config import UPLOAD_DIR
 from service.video.BaseVideoAnalyzer import BaseVideoAnalyzer
 
 
@@ -149,54 +147,6 @@ from service.video.BaseVideoAnalyzer import BaseVideoAnalyzer
 #     output.close()
 
 
-def split_video(video_uuid):
-    video_path = UPLOAD_DIR + video_uuid + ".mp4"
-    frames_dir = UPLOAD_DIR + "frames/" + video_uuid + "/"
-
-    if os.path.exists(frames_dir):  # Already split
-        return
-
-    os.makedirs(frames_dir)
-
-    cap = cv2.VideoCapture(video_path)
-    FPS = cap.get(5)  # 5 means fps
-    logging.debug("Splitting video %s, FPS: %s" % (video_path, FPS))
-
-    c = 1
-    sampling_fps = 1
-
-    while True:
-        ret, frame = cap.read()
-        if ret:
-            if c % sampling_fps == 0:
-                target_path = frames_dir + str(c) + ".jpg"
-                cv2.imwrite(target_path, frame)
-                logging.debug("Writing frame %s, %s" % (c, target_path))
-            c += 1
-        else:
-            break
-
-    cap.release()
-
-
-from service.aiutil import torch_openpose
-from service.aiutil.util import draw_bodypose
-
-
-def run_openpose_for_frames(video_uuid):
-    tp = torch_openpose('body_25')
-
-    frames_dir = UPLOAD_DIR + "frames/" + video_uuid + "/"
-    for i in os.listdir(frames_dir):
-        if i.endswith(".jpg"):
-            oriImg = cv2.imread(frames_dir + i)
-            poses = tp(oriImg)
-            np.save(frames_dir + i + ".npy", poses)
-
-            canvas = draw_bodypose(oriImg, poses, 'body_25')
-            cv2.imwrite(frames_dir + i.replace(".jpg", ".result"), canvas)
-
-
 class Video3DAnalyzer(BaseVideoAnalyzer):
     """Perform inference on a single video"""