-
Notifications
You must be signed in to change notification settings - Fork 0
/
grabcut.py
101 lines (79 loc) · 3.5 KB
/
grabcut.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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Dec 23 15:53:25 2018
@author: ibrahim
"""
import numpy as np
from enum import Enum
Trimap = Enum('Trimap', 'TrimapBackground TrimapUnknown TrimapForeground')
Matte = Enum('Matte', 'MatteBackground MatteForeground')
from commonfunctions import *
img = io.imread('images/house 2-small.jpg')
from sklearn.mixture import GaussianMixture
GMM = GaussianMixture(n_components=5)
img.shape
GMM.fit(img.reshape((img.shape[0]*img.shape[1], 3)))
#print(GMM.weights_.shape)
#print(GMM.weights_)
#print("tttt")
#print(GMM.means_.shape)
#print(GMM.means_)
#print("yyyy")
#print(GMM.covariances_.shape)
#print(GMM.covariances_)
x,y = 200,200
width,height = 100,100
# initialize all as BG
trimap = np.full((img.shape[0], img.shape[1]), Trimap.TrimapBackground)
trimap[y:y+height, x:x+width] = Trimap.TrimapUnknown
#print(trimap)
matte = np.full((img.shape[0],img.shape[1]), Matte.MatteBackground)
matte[trimap == Trimap.TrimapUnknown] = Matte.MatteForeground
foreground_pixels = img[matte == Matte.MatteForeground]
background_pixels = img[matte == Matte.MatteBackground]
#BG_GMM = GaussianMixture(n_components=5).fit(background_pixels)
#FG_GMM = GaussianMixture(n_components=5).fit(foreground_pixels)
#
#GMM_components = np.empty_like(matte).astype(int)
#
#for col in range(img.shape[0]):
# for row in range(img.shape[1]):
# pixel = img[col,row]
# if matte[col,row] == Matte.MatteForeground:
# proba = FG_GMM.predict_proba([pixel])
## D = [-np.log10(FG_GMM.weights_[i] * (1/np.sqrt(np.linalg.det(FG_GMM.covariances_[i]))) * np.exp(-0.5*(pixel-FG_GMM.means_[i]).dot(np.linalg.inv(FG_GMM.covariances_[i])).dot(pixel - FG_GMM.means_[i]))) for i in range(5)]
# else:
# proba = BG_GMM.predict_proba([pixel])
## D = [-np.log10(BG_GMM.weights_[i] * (1/np.sqrt(np.linalg.det(BG_GMM.covariances_[i]))) * np.exp(-0.5*(pixel-BG_GMM.means_[i]).dot(np.linalg.inv(BG_GMM.covariances_[i])).dot(pixel - BG_GMM.means_[i]))) for i in range(5)]
## if(proba.argmax() != np.array(D).argmin()):
## print("proba", proba)
## print("D", D)
# GMM_components[col,row] = proba.argmax()
all_GMM_components = np.empty_like(matte).astype(int)
for col in range(img.shape[0]):
for row in range(img.shape[1]):
pixel = img[col,row]
proba = GMM.predict_proba([pixel])
all_GMM_components[col,row] = proba.argmax()
gmm_mask = [all_GMM_components == i for i in range(5)]
fore_mask = matte == Matte.MatteForeground
#foreground_pixels = [ img[matte == Matte.MatteForeground and all_GMM_components == i] for i in range(5)]
#background_pixels = [img[matte == Matte.MatteBackground and all_GMM_components == i] for i in range(5)]
fore_pixels = [img[fore_mask * gmm_mask[i]] for i in range(5)]
back_pixels = [img[np.invert(fore_mask) * gmm_mask[i]] for i in range(5)]
#fore_GMM = [GaussianMixture(n_components=1).fit(fore_pixel) for fore_pixel in fore_pixels if fore_pixel.shape[0] != 0]
#back_GMM = [GaussianMixture(n_components=1).fit(back_pixel) for back_pixel in back_pixels if back_pixel.shape[0] != 0]
fore_GMM = []
back_GMM = []
for fore_pixel in fore_pixels:
g = None
if (fore_pixel.shape[0] != 0):
g = GaussianMixture(n_components=1).fit(fore_pixel)
fore_GMM.append(g)
for back_pixel in back_pixels:
g = None
if (back_pixel.shape[0] != 0):
g = GaussianMixture(n_components=1).fit(back_pixel)
back_GMM.append(g)
# iterate over and solve min-cut/max-flow for graph