import numpy as np
import os
import torch.nn as nn

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import cv2
import torch

from . import util
from .body import Body

remote_model_path = "https://huggingface.co/TencentARC/T2I-Adapter/blob/main/third-party-models/body_pose_model.pth"


class OpenposeInference(nn.Module):

    def __init__(self):
        super().__init__()
        body_modelpath = os.path.join('models', "body_pose_model.pth")

        if not os.path.exists(body_modelpath):
            from basicsr.utils.download_util import load_file_from_url
            load_file_from_url(remote_model_path, model_dir='models')

        self.body_estimation = Body(body_modelpath)

    def forward(self, x):
        x = x[:, :, ::-1].copy()
        with torch.no_grad():
            candidate, subset = self.body_estimation(x)
            canvas = np.zeros_like(x)
            canvas = util.draw_bodypose(canvas, candidate, subset)
            canvas = cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR)
        return canvas