PoseCorrection/backend_refactor/service/photo_analyzer.py

import copy
from typing import Callable

import cv2

from config import UPLOAD_DIR, FILE_URL
from service.ai.photo.standing_photo import *
from service.ai.util.torch_openpose import torch_openpose
from service.ai.util.util import draw_bodypose


class BasePhotoAnalyzer:
    def __init__(self):
        pass


class StandingPhotoAnalyzer(BasePhotoAnalyzer):

    def __init__(self, data: dict):
        super().__init__()
        assert len(data) == 3, "Invalid arguments: expect 2 photo uuids and height."
        self.front = data['front']
        self.right = data['right']
        self.height = data['height']

    def analyze(self, callback: Callable) -> None:
        try:
            callback("RUNNING", progress=0)
            photos = []
            npys = []

            for index, i in enumerate((self.front, self.right)):
                test_image = UPLOAD_DIR + '{}.jpg'.format(i)

                ori_img = cv2.imread(test_image)  # B,G,R order

                tp = torch_openpose('body_25')
                poses = tp(ori_img)

                if index == 0:
                    npys.append(analyse_npy_front(poses, self.height))
                else:
                    npys.append(analyse_npy_side(poses, self.height))

                canvas = copy.deepcopy(ori_img)
                canvas = draw_bodypose(canvas, poses, 'body_25')

                cv2.imwrite(UPLOAD_DIR + 'result/{}.result.jpg'.format(i), canvas)
                photos.append(FILE_URL + '{}.result.jpg'.format(i))

            callback("FINISHED", result={"data": {"front": {"file": photos[0]}, "right": {"file": photos[1]}},
                                         "advice": {"front": npys[0], "right": npys[1], "points": poses}}, )
        except Exception as e:
            callback("ERROR", error=e)
            raise e