import numpy as np
import cv2

# 食道の映像を前処理することが多いような気がするので作ったクラス
# Undistort用のparamを同ディレクトリに配置する必要あり
class EsoMovieConverter:

    def __init__(self):
        self.dist_coeffs = np.load('./params/dist_coeffs.npy')
        self.intrinsics_scaled = np.array([[221.8766, 0, 232.2143], [0, 217.4069, 173.3776], [0, 0, 1]])
        self.kernel = np.ones((9, 9), np.uint8)

    def __call__(self, eso_frame):
        eso_frame = eso_frame[32:989, 323:1599, :]
        # eliminate interrace
        eso_frame = eso_frame[::2, ::2, :]
        eso_frame = cv2.resize(eso_frame, (480, 352), interpolation=cv2.INTER_LINEAR)
        eso_frame = cv2.undistort(eso_frame, self.intrinsics_scaled, self.dist_coeffs)

        # yuv_frame = cv2.cvtColor(eso_frame, cv2.COLOR_BGR2YUV)
        # mask = cv2.dilate(np.where(200 < yuv_frame[:, :, 0], 255, 0).astype(yuv_frame.dtype), self.kernel)
        # eso_frame = cv2.inpaint(eso_frame, mask, 5, cv2.INPAINT_TELEA)
        return eso_frame