PoseCorrection/backend_refactor/service/video_analyzer.py

import logging
import os.path
import pickle
from typing import Callable

import cv2

from config import UPLOAD_DIR, FILE_URL
from service.ai.util.torch_openpose import torch_openpose
from service.ai.util.util import draw_bodypose
from service.ai.video.classify.balancebeam import classify_balancebeam
from service.ai.video.classify.gaotaitui import classify_gaotaitui
from service.ai.video.classify.juanfu import classify_juanfu
from service.ai.video.classify.jump import classify_jump
from service.ai.video.classify.jumpwithboth import classify_jumpwithboth
from service.ai.video.classify.pingban import classify_pingban
from service.ai.video.classify.shendun import classify_shendun
from service.ai.video.classify.sitforward import classify_sitforward
from service.ai.video.classify.standing import classify_standing
from service.ai.video.classify.tennisthrow import classify_tennisthrow
from service.ai.video.classify.turnaround import classify_turnaround
from service.ai.video.inner_analyze import analyse_npy_side_jump, analyse_npy_side_juanfu, analyse_npy_side_shendun, \
    analyse_npy_side_gaotaitui
from service.ai.video.video_score.balancebeam import cal_balancebeam
from service.ai.video.video_score.gaotaitui import ignore_data_gaotaitui
from service.ai.video.video_score.juanfu import ignore_data_juanfu
from service.ai.video.video_score.jump import cal_jump
from service.ai.video.video_score.jumpwithboth import cal_jumpwithboth
from service.ai.video.video_score.pingban import cal_pingban
from service.ai.video.video_score.shendun import ignore_data_shendun
from service.ai.video.video_score.sitforward import cal_sitforward
from service.ai.video.video_score.standing import cal_standing
from service.ai.video.video_score.tennisthrow import cal_tennisthrow
from service.ai.video.video_score.turnaround import cal_turnaround


def run_openpose_for_npy_and_video(video_uuid, callback: Callable):
    video_path = UPLOAD_DIR + video_uuid + ".mp4"

    cap = cv2.VideoCapture(video_path)
    FPS = cap.get(cv2.CAP_PROP_FPS)
    total_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
    out = None
    logging.debug("Splitting video %s, FPS: %s" % (video_path, FPS))

    c = 1
    sampling_fps = 1

    tp = torch_openpose('body_25')  # AI Model

    results = {}

    while True:
        ret, frame = cap.read()
        if ret:
            if c % sampling_fps == 0:
                logging.debug("Running openpose for %s %s" % (video_uuid, c))

                poses = tp(frame)

                results[str(c)] = poses  # this is evil, but the AI code expect it to be so.

                canvas = draw_bodypose(frame, poses, 'body_25')
                cv2.imwrite(UPLOAD_DIR + "result/" + video_uuid + "." + str(c) + ".jpg", canvas)

                if out is None:
                    out = cv2.VideoWriter(UPLOAD_DIR + "result/" + video_uuid + ".result.mp4",
                                          cv2.VideoWriter_fourcc(*"vp90"), FPS,
                                          (canvas.shape[1], canvas.shape[0]))  # So fucking stupid
                out.write(canvas)

            c += 1
            callback("RUNNING", progress=min(99, int(100 * c / (total_frames / sampling_fps))))
        else:
            break

    cap.release()
    out.release()

    return results


class BaseVideoAnalyzer:
    def __init__(self, data: dict):
        self.video_uuid = data['video_uuid']
        self.results = None
        self.points = None

    def _do_analyze(self):
        assert False, "Internal error: You should override this with a subclass!"

    def analyze(self, callback: Callable):
        try:
            callback("RUNNING", progress=0)

            PICKLE_DUMP_PATH = UPLOAD_DIR + self.video_uuid + ".dump"

            if not os.path.exists(PICKLE_DUMP_PATH):  # analyze
                self.results = run_openpose_for_npy_and_video(self.video_uuid, callback)
                self.points = self.results
                pickle.dump(self.results, open(PICKLE_DUMP_PATH, "wb"))
            else:  # read cache
                self.results = pickle.load(open(PICKLE_DUMP_PATH, "rb"))

            result = {**self._do_analyze(), "file": FILE_URL + self.video_uuid + ".result.mp4"}
            # inject file url here
            try:
                for i in result['data']:
                    result['data'][i]['file'] = FILE_URL + self.video_uuid + "." + result['data'][i]['index'] + ".jpg"
                result['data']['points']=self.points
            except Exception as e:
                logging.warning("file index in data not found, %s" % e)

            callback("FINISHED", result=result)
        except Exception as e:
            callback("ERROR", error=e)
            raise e


class JumpVideoAnalyzer(BaseVideoAnalyzer):
    def _do_analyze(self):
        self.results = {k: analyse_npy_side_jump(v) for k, v in self.results.items()}
        self.results = {k: v for k, v in self.results.items() if v != {}}
        raw = cal_jump(self.results)[1]
        return classify_jump(raw)


class PingbanVideoAnalyzer(BaseVideoAnalyzer):
    def _do_analyze(self):
        self.results = {k: analyse_npy_side_jump(v) for k, v in self.results.items()}
        self.results = {k: v for k, v in self.results.items() if v != {}}
        raw = cal_pingban(self.results)[1]
        return classify_pingban(raw)


class JuanfuVideoAnalyzer(BaseVideoAnalyzer):
    def _do_analyze(self):
        self.results = {k: analyse_npy_side_juanfu(v) for k, v in self.results.items()}
        raw = ignore_data_juanfu(self.results)[1]
        return classify_juanfu(raw)


class ShendunVideoAnalyzer(BaseVideoAnalyzer):
    def _do_analyze(self):
        self.results = {k: analyse_npy_side_shendun(v) for k, v in self.results.items()}
        raw = ignore_data_shendun(self.results)[1]
        return classify_shendun(raw)


class GaotaituiVideoAnalyzer(BaseVideoAnalyzer):
    def _do_analyze(self):
        self.results = {k: analyse_npy_side_gaotaitui(v) for k, v in self.results.items()}
        raw = ignore_data_gaotaitui(self.results)[1]
        return classify_gaotaitui(raw)


class TurnaroundVideoAnalyzer(BaseVideoAnalyzer):
    def _do_analyze(self):
        self.results = {k: analyse_npy_side_jump(v) for k, v in self.results.items()}
        raw = cal_turnaround(self.results)[1]
        return classify_turnaround(raw)


class TennisthrowVideoAnalyzer(BaseVideoAnalyzer):
    def _do_analyze(self):
        self.results = {k: analyse_npy_side_jump(v) for k, v in self.results.items()}
        raw = cal_tennisthrow(self.results)[1]
        return classify_tennisthrow(raw)


class JumpwithbothVideoAnalyzer(BaseVideoAnalyzer):
    def _do_analyze(self):
        self.results = {k: analyse_npy_side_jump(v) for k, v in self.results.items()}
        raw = cal_jumpwithboth(self.results)[1]
        return classify_jumpwithboth(raw)


class SitforwardVideoAnalyzer(BaseVideoAnalyzer):
    def _do_analyze(self):
        self.results = {k: analyse_npy_side_jump(v) for k, v in self.results.items()}
        raw = cal_sitforward(self.results)[1]
        return classify_sitforward(raw)


class BalancebeamVideoAnalyzer(BaseVideoAnalyzer):
    def _do_analyze(self):
        self.results = {k: analyse_npy_side_jump(v) for k, v in self.results.items()}
        raw = cal_balancebeam(self.results)[1]
        return classify_balancebeam(raw)


class StandingVideoAnalyzer(BaseVideoAnalyzer):
    def _do_analyze(self):
        self.results = {k: analyse_npy_side_jump(v) for k, v in self.results.items()}
        raw = cal_standing(self.results)[1]
        return classify_standing(raw)