from mhr.common import *

class PIL2CV(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, img):
        return cv.cvtColor(np.asarray(img), cv.COLOR_RGB2BGR)

class CV2PIL(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, img):
        return Image.fromarray(cv.cvtColor(img, cv.COLOR_BGR2RGB))

class Tensor2CV(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, img):
        if len(img.shape) == 4:
            img = img.squeeze(0)
        img = img.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).type(torch.uint8).numpy()
        img = cv.cvtColor(img, cv.COLOR_RGB2BGR)
        return img


class CvtColor(torch.nn.Module):
    def __init__(self, cvt):
        super().__init__()
        self._cvt = cvt
    def forward(self, img):
        return cv.cvtColor(img, self._cvt)

class GaussianBlur(torch.nn.Module):
    def __init__(self, kernel, sigma):
        super().__init__()
        self._kernel = kernel
        self._sigma = sigma
    def forward(self, img):
        return cv.GaussianBlur(img, self._kernel, self._sigma)

class EqualizeHist(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, img):
        return cv.equalizeHist(img)

class SobelBitwiseOrXY(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, img):
        sobelx = cv.Sobel(img, cv.CV_64F, 1, 0)
        sobely = cv.Sobel(img, cv.CV_64F, 0, 1)
        sobelx = cv.convertScaleAbs(sobelx)
        sobely = cv.convertScaleAbs(sobely)
        return cv.bitwise_or(sobelx, sobely)

class Threshold(torch.nn.Module):
    def __init__(self, thresh, maxval, tt):
        super().__init__()
        self._thresh = thresh
        self._maxval = maxval
        self._tt = tt
    def forward(self, img):
        _, tmp = cv.threshold(img, self._thresh, self._maxval, self._tt)
        return tmp

class Cut(torch.nn.Module):
    def __init__(self, point, offsetx, offsety):
        super().__init__()
        self._p = point
        self._offsetx = offsetx
        self._offsety = offsety
    def forward(self, img):
        p = self._p
        return img[p[0]:p[0]+self._offsetx, p[1]:p[1]+self._offsety ]



class TensorCut(torch.nn.Module):
    def __init__(self, point, offsetx, offsety):
        super().__init__()
        self._p = point
        self._offsetx = offsetx
        self._offsety = offsety
    def forward(self, img):
        p = self._p
        return img[:, p[0]:p[0]+self._offsetx, p[1]:p[1]+self._offsety ]