import numpy as np
from PIL import Image
import paddle.inference as paddle_infer
import paddle.vision.transforms as T


class corner_finder():
    def __init__(self, model, params):
        config = paddle_infer.Config(model, params)
        self.predictor = paddle_infer.create_predictor(config)

    def get_location(self, img, retainFactor=0.85):
        trans = T.ToTensor()
        resi = T.Resize([32, 32])

        ans_x = 0.0
        ans_y = 0.0

        o_img = np.copy(img)

        y = [0, 0]
        x_start = 0
        y_start = 0
        up_scale_factor = (img.shape[1], img.shape[0])

        myImage = np.copy(o_img)

        CROP_FRAC = retainFactor
        while (myImage.shape[0] > 10 and myImage.shape[1] > 10):
            img_temp = Image.fromarray(myImage)
            img_temp = trans(img_temp)
            img_temp = resi(img_temp)
            img_temp = np.expand_dims(img_temp, 0)
            input_names = self.predictor.get_input_names()
            input_handle = self.predictor.get_input_handle(input_names[0])

            input_handle.reshape([1, 3, 32, 32])
            input_handle.copy_from_cpu(img_temp)

            # 运行predictor
            self.predictor.run()

            # 获取输出
            output_names = self.predictor.get_output_names()
            output_handle = self.predictor.get_output_handle(output_names[0])
            output_data = output_handle.copy_to_cpu()  # numpy.ndarray类型
            output_data = output_data[0]
            model_prediction = np.array(output_data)

            response_up = model_prediction

            response_up = response_up * up_scale_factor
            y = response_up + (x_start, y_start)
            x_loc = int(y[0])
            y_loc = int(y[1])

            if x_loc > myImage.shape[1] / 2:
                start_x = min(x_loc + int(round(myImage.shape[1] * CROP_FRAC / 2)), myImage.shape[1]) - int(round(
                    myImage.shape[1] * CROP_FRAC))
            else:
                start_x = max(x_loc - int(myImage.shape[1] * CROP_FRAC / 2), 0)
            if y_loc > myImage.shape[0] / 2:
                start_y = min(y_loc + int(myImage.shape[0] * CROP_FRAC / 2), myImage.shape[0]) - int(
                    myImage.shape[0] * CROP_FRAC)
            else:
                start_y = max(y_loc - int(myImage.shape[0] * CROP_FRAC / 2), 0)

            ans_x += start_x
            ans_y += start_y

            myImage = myImage[start_y:start_y + int(myImage.shape[0] * CROP_FRAC),
                        start_x:start_x + int(myImage.shape[1] * CROP_FRAC)]
            img = img[start_y:start_y + int(img.shape[0] * CROP_FRAC),
                    start_x:start_x + int(img.shape[1] * CROP_FRAC)]
            up_scale_factor = (img.shape[1], img.shape[0])

        ans_x += y[0]
        ans_y += y[1]
        return (int(round(ans_x)), int(round(ans_y)))


if __name__ == "__main__":
    pass