# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/08_inference.ipynb (unless otherwise specified).

__all__ = ['KeypointInference']

# Cell
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import os
import json
import numpy as np
import cv2
import torch
from skimage import io
import matplotlib.pyplot as plt
import random
from tqdm import tqdm

# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg


from fastscript import call_parse
from fastscript import *
from fastcore.transform import store_attr

from .perspective_correction import four_point_transform
from .utils import load_yaml

# Cell

class KeypointInference():
    def __init__(self, infer_cfg: str, kp_type: str):
        store_attr(self, "infer_cfg")
        self.config = load_yaml(infer_cfg)

        ROOT = self.config['ROOT']
        model_cfg = os.path.join(ROOT, self.config['CONFIGS'][kp_type])
        model_path = os.path.join(ROOT, self.config['WEIGHTS'][kp_type])

        self.cfg = get_cfg()
        self.cfg.merge_from_file(model_zoo.get_config_file(self.config['BASE_CONFIGS']['keypoint']))
        self.cfg.merge_from_file(model_cfg)

        if not torch.cuda.is_available(): self.cfg.MODEL.DEVICE = 'cpu';
        self.cfg.MODEL.WEIGHTS = model_path
        self.predictor = DefaultPredictor(self.cfg)

    def predict_keypoints(self, img_loc, visualize=False):
        img = io.imread(img_loc) if isinstance(img_loc, str) else img_loc
        outputs = self.predictor(img)
        keypoints = outputs["instances"].pred_keypoints.to('cpu').numpy().astype(np.int32)[:,:,:2]
        #return if no keypoints are detected
        if len(keypoints)<1:
            return keypoints, False
        vis_img = self.vizualize_keypoints(img.copy(),keypoints[0])
        if visualize:
            plt.imshow(vis_img)
        return keypoints, vis_img

    def vizualize_keypoints(self,img, pred_kp,figure_size=(15,9),isClosed = True,thickness = 7,color = (255, 0, 0)):

        for j in range(-1,3):
            x1,y1,x2,y2 = pred_kp[j][0],pred_kp[j][1],pred_kp[j+1][0], pred_kp[j+1][1]
            cv2.circle(img,(x1, y1), 9, (0,255,0), -1)
            cv2.line(img,(x1,y1),(x2,y2),color,thickness)
        return img