human/render/video_render.py

#encoding = utf8
import copy
import logging
import time

import cv2
import numpy as np

from .base_render import BaseRender


def img_warp_back_inv_m(img, img_to, inv_m):
    h_up, w_up, c = img_to.shape
    t = time.perf_counter()
    mask = np.ones_like(img).astype(np.float32)
    inv_mask = cv2.warpAffine(mask, inv_m, (w_up, h_up))
    inv_img = cv2.warpAffine(img, inv_m, (w_up, h_up))
    mask_indices = inv_mask == 1
    print(f'time1: {time.perf_counter() - t}')
    if 4 == c:
        t = time.perf_counter()
        img_to[:, :, :3][mask_indices] = inv_img[mask_indices]
        print(f'time2: {time.perf_counter() - t}')
    else:
        img_to[inv_mask == 1] = inv_img[inv_mask == 1]
    return img_to


class VideoRender(BaseRender):
    def __init__(self, play_clock, context, human_render):
        super().__init__(play_clock, context, 'Video')
        self._human_render = human_render
        self.index = 0

    def render(self, frame, ps):
        res_frame, idx, type_ = frame

        if type_ == 0:
            combine_frame = self._context.frame_list_cycle[idx]
        else:
            bbox = self._context.coord_list_cycle[idx]
            combine_frame = copy.deepcopy(self._context.frame_list_cycle[idx])
            af = self._context.align_frames[idx]
            inv_m = self._context.inv_m_frames[idx]
            y1, y2, x1, x2 = bbox
            try:
                t = time.perf_counter()
                res_frame = cv2.resize(res_frame.astype(np.uint8), (x2 - x1, y2 - y1))
                af[y1:y2, x1:x2] = res_frame
                combine_frame = img_warp_back_inv_m(af, combine_frame, inv_m)
                print(time.perf_counter() - t)
            except Exception as e:
                logging.error(f'resize error:{e}')
                return

        image = combine_frame
        # image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)
        if self._human_render is not None:
            self._human_render.put_image(image)