-
Notifications
You must be signed in to change notification settings - Fork 0
/
IHMeans.py
76 lines (61 loc) · 2.09 KB
/
IHMeans.py
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
import cv2
import numpy as np
import matplotlib.pyplot as plt
import random
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--iter', dest='iter', type=int, default=3, help='# of epoch')
parser.add_argument('--input', dest='input', type=str)
args = parser.parse_args()
img_noise = cv2.imread('./image/%s.png' % args.input)
b, g, r = cv2.split(img_noise)
img_noise = cv2.merge([r, g, b])
mask = []
for row in range(img_noise.shape[0]):
line_mask = []
for col in range(img_noise.shape[1]):
if img_noise[row][col][0] != 0 or img_noise[row][col][1] != 0 or img_noise[row][col][2] != 0:
line_mask.append(1)
else:
line_mask.append(0)
mask.append(line_mask)
def IHMeanMask(image):
img_mask = image.copy()
for row in range(image.shape[0]):
for col in range(image.shape[1]):
if mask[row][col] == 1:
img_mask[row][col] = img_noise[row][col]
else:
pass
return img_mask
def IHMeanOperator(roi, q):
roi = roi.astype(np.float64)
return np.mean((roi)**(q+1))/np.mean((roi)**(q))
def IHMeanFilter(image, q):
image_filter = np.zeros(image.shape)
image = cv2.copyMakeBorder(image, 1, 1, 1, 1, cv2.BORDER_DEFAULT)
for i in range(1, image.shape[0]-1):
for j in range(1, image.shape[1]-1):
image_filter[i-1, j-1] = IHMeanOperator(image[i-1:i+2, j-1:j+2], q)
image_filter = (image_filter-np.min(image))*(255/np.max(image))
return image_filter.astype(np.uint8)
def IHMean(image, q):
r, g, b = cv2.split(image)
r = IHMeanFilter(r, q)
g = IHMeanFilter(g, q)
b = IHMeanFilter(b, q)
return cv2.merge([r, g, b])
def IHMeanEpoch(image):
img_IHMean = IHMeanMask(image)
img_IHMean = IHMean(img_IHMean, 0.1)
return img_IHMean
print('[*] Start!')
print('[*] Iter 0')
img_IHMean = IHMean(img_noise, 0.1)
for i in range(args.iter):
print('[*] Iter %d' % (i+1))
img_IHMean = IHMeanEpoch(img_IHMean)
r, g, b = cv2.split(img_IHMean)
img_IHMean = cv2.merge([b, g, r])
cv2.imwrite("result.png", img_IHMean)
print('[*] Finish!')