#!/usr/bin/env python
# encoding: utf-8
"""
@author: simin. Gong
@license: (C) Copyright 2013-2023 Spacemit Corporation Limited.
@contact: simin.gong@spacemit.com
@file: rtm_pose_node.py
@date: 2024/12/17 14:18 
@desc:
"""

import vp_module as vp
from infer_module import RTMPost as rtmpost
import numpy as np

class RTMPostNode(vp.vp_node):
    def __init__(self, node_name, model_path, callback, mode = vp.QueueMode.DropOldest):
        super().__init__(node_name, mode)
        self.initialized()
        self.engin = rtmpost.create_infer(model_path)
        self.callback = callback

    def __del__(self):
        self.gate_open()
        self.deinitialized()

    def handle_frame_meta(self, meta):
        height, width = meta.frame.shape[:2]
        targets = meta.any_targets.get()
        results = []
        for box in targets:
            for Box in box.boxs:
                keypoints_data = []
                if Box.label == 0:
                    x1, y1, x2, y2 = int(Box.rect.x * width), int(Box.rect.y * height), int(
                        (Box.rect.x + Box.rect.width) * width), int((Box.rect.y + Box.rect.height) * height)
                    x1 = max(0, x1)
                    y1 = max(0, y1)
                    x2 = min(meta.frame.shape[1] - 1, x2)
                    y2 = min(meta.frame.shape[0] - 1, y2)
                    person_roi = meta.frame[y1:y2, x1:x2]
                    pose_result = self.engin.commit(person_roi.copy()).get()
                    for keypoint in pose_result:
                        keypoints = keypoint.keypoints
                        for kp in keypoints:
                            # kp_x = np.clip(kp.x, x1, x2)
                            # kp_y = np.clip(kp.y, y1, y2)
                            kp_x = float((kp.x + x1) / width)
                            kp_y = float((kp.y + y1) / height)
                            kp_x = np.clip(kp_x, 0, 1)
                            kp_y = np.clip(kp_y, 0, 1)
                            kp_score = kp.score
                            keypoints_data.append({
                                "x": kp_x,
                                "y": kp_y,
                                "score": kp_score
                            })
                person_data = {"keypoints": keypoints_data}
                results.append(person_data)
        # json_str = json.dumps(results, ensure_ascii=False, indent=4)
        # print(json_str)
        if self.callback:
            self.callback(None, results)
        return super().handle_frame_meta(meta)
