fix 'get_avg_color' and implement score function
This commit is contained in:
parent
5a51510209
commit
2f98d26842
1 changed files with 30 additions and 3 deletions
33
classify.py
33
classify.py
|
@ -4,6 +4,10 @@ import h5py as h5
|
|||
|
||||
db = None
|
||||
|
||||
def init_train():
|
||||
''' Default init based on the train set `train.h5` '''
|
||||
init('train.h5')
|
||||
|
||||
def init(path):
|
||||
''' initializes the database, must be called before any use '''
|
||||
global db
|
||||
|
@ -51,18 +55,41 @@ def get_avg_color(img, mask):
|
|||
gets avg color from an image that is underneath a mask,
|
||||
img and mask needs to be of same size(in x,y) but img can
|
||||
have any third dimension size it want(usually 3 for rgb or 1 for grayscale)
|
||||
|
||||
mask needs to be of shape(img.width, img.height, 1)
|
||||
'''
|
||||
sx, sy, sw = img.shape
|
||||
mx, my = mask.shape
|
||||
if sx != mx or sy != my:
|
||||
print('Image and mask shape doesnt match!')
|
||||
print('Image and mask size doesnt match!')
|
||||
return None
|
||||
avg = np.array(sw, dtype=np.float32)
|
||||
avg = np.zeros(sw, dtype=np.float32)
|
||||
count = 0.0
|
||||
for x in range(sx):
|
||||
for y in range(sy):
|
||||
m = mas[x, y]
|
||||
m = mask[x, y]
|
||||
avg += img[x, y] * m
|
||||
count += m
|
||||
avg /= count
|
||||
return avg
|
||||
|
||||
def calc_score(img, mask, avg_color):
|
||||
'''
|
||||
Calculates the score for each mask with each color
|
||||
'''
|
||||
sx, sy, sw = img.shape
|
||||
mx, my = mask.shape
|
||||
if sx != mx or sy != my:
|
||||
print('Image and mask size doesnt match!')
|
||||
return 0.0
|
||||
if sw != avg_color.shape[0]:
|
||||
print('Image width doesnt match color width!')
|
||||
return 0.0
|
||||
score = 0.0
|
||||
for x in range(sx):
|
||||
for y in range(sy):
|
||||
m = 0.5 - mask[x, y]
|
||||
diff = img[x, y] - avg_color
|
||||
mag = np.sqrt(diff.dot(diff)) # calculate magnitude
|
||||
score += mag * m
|
||||
return score
|
||||
|
|
Loading…
Reference in a new issue