Spaces:
Sleeping
Sleeping
File size: 2,601 Bytes
a22da3a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
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 ] |