├── .gitignore ├── GraphMaker.py ├── Imgs ├── Depth │ └── 8_08-40-06_Depth.png └── RGB │ └── 8_08-40-06.jpg ├── README.md ├── Results └── .gitignore ├── doc ├── gui1.png ├── label.png ├── segment_cues.jpg └── top.png └── main.py /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | -------------------------------------------------------------------------------- /GraphMaker.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import maxflow 4 | from skimage.segmentation import slic 5 | from skimage.segmentation import mark_boundaries 6 | from skimage import img_as_ubyte 7 | from dijkstar import Graph, find_path 8 | import os 9 | 10 | import datetime 11 | 12 | 13 | class GraphMaker: 14 | foreground = 1 15 | background = 0 16 | 17 | seeds = 0 18 | segmented = 1 19 | 20 | default = 0.5 21 | MAXIMUM = 1000000000 22 | 23 | ImageName = "8_08-40-06" 24 | 25 | ns = 1 26 | ###parameters### 27 | lamda = 0.2 28 | NORMAL = True # Consider normal vector or not? 29 | 30 | def __init__(self): 31 | self.depth = None 32 | self.image = None 33 | self.superpixel_image = None 34 | self.overlay = None 35 | self.seed_overlay = None 36 | self.segment_overlay = None 37 | self.load_image(os.path.join('Imgs', 'RGB', self.ImageName+'.jpg'), 38 | os.path.join('Imgs', 'Depth', self.ImageName+'_Depth.png')) 39 | self.background_seeds = [] 40 | self.foreground_seeds = [] 41 | self.foreground_superseeds = [] 42 | self.background_superseeds = [] 43 | self.nodes = [] 44 | self.edges = [] 45 | self.current_overlay = self.seeds 46 | 47 | self.ave_LAB = None 48 | self.ave_normal = None 49 | self.ave_depth = None 50 | self.confident_map = None 51 | self.cue_selection = None 52 | self.superpixel_segment = None 53 | self.super_edge = None 54 | 55 | self.normal_map = None 56 | self.LAB_map = None 57 | self.depth_map = None 58 | self.cfd_LAB = None 59 | self.cfd_normal = None 60 | self.cfd_depth = None 61 | self.super_label = None # label of superpixels，0~5 for 62 | # 0:LAB FG 1:LAB BG 2:depth FG 63 | # 3:depth BG 4:normal FG 5:normal BG 64 | 65 | def computeSigma(self): 66 | 67 | self.sigma1 = 0 68 | self.sigma2 = 0 69 | self.sigma3 = 0 70 | 71 | for i in range(self.n_seg): 72 | for j in self.super_edge[i]: 73 | if j > i: 74 | if self.sigma1 < self.eu_dis(self.ave_LAB[i], self.ave_LAB[j]): 75 | self.sigma1 = self.eu_dis(self.ave_LAB[i], self.ave_LAB[j]) 76 | if self.sigma2 < np.abs(self.ave_depth[i] - self.ave_depth[j]): 77 | self.sigma2 = np.abs(self.ave_depth[i] - self.ave_depth[j]) 78 | if self.NORMAL and self.sigma3 < (1 - self.cosine_sim(self.ave_normal[i], self.ave_normal[j])): 79 | self.sigma3 = (1 - self.cosine_sim(self.ave_normal[i], self.ave_normal[j])) 80 | 81 | self.sigma1 = self.sigma1 ** 2 * 1 82 | self.sigma2 = self.sigma2 ** 2 * 1 83 | self.sigma3 = self.sigma3 ** 2 * 0.5 84 | print("sigma1 = " + str(self.sigma1)) 85 | print("sigma2 = " + str(self.sigma2)) 86 | print("sigma3 = " + str(self.sigma3)) 87 | 88 | def load_image(self, filename, depth_filename): 89 | self.image = cv2.imread(filename) 90 | self.depth = cv2.imread(depth_filename, cv2.IMREAD_ANYDEPTH) 91 | self.superpixel_image = self.image.copy() 92 | self.seed_overlay = np.zeros_like(self.image) 93 | self.segment_overlay = np.zeros_like(self.image) 94 | 95 | def add_seed(self, x, y, type): 96 | if self.image is None: 97 | print('Please load an image before adding seeds.') 98 | if type == self.background: 99 | if not self.background_seeds.__contains__((x, y)): 100 | self.background_seeds.append((x, y)) 101 | cv2.rectangle(self.seed_overlay, (x - 1, y - 1), (x + 1, y + 1), (0, 0, 255), -1) 102 | elif type == self.foreground: 103 | if not self.foreground_seeds.__contains__((x, y)): 104 | self.foreground_seeds.append((x, y)) 105 | cv2.rectangle(self.seed_overlay, (x - 1, y - 1), (x + 1, y + 1), (0, 255, 0), -1) 106 | 107 | def clear_seeds(self): 108 | self.background_seeds = [] 109 | self.foreground_seeds = [] 110 | self.background_superseeds = [] 111 | self.foreground_superseeds = [] 112 | self.seed_overlay = np.zeros_like(self.seed_overlay) 113 | 114 | def get_image_with_overlay(self, overlayNumber): 115 | if overlayNumber == self.seeds: 116 | return cv2.addWeighted(self.image, 0.9, self.seed_overlay, 0.4, 0.1) 117 | else: 118 | return cv2.addWeighted(self.image, 0.3, self.segment_overlay, 0.7, 0.1) 119 | 120 | def create_graph(self): 121 | starttime = datetime.datetime.now() 122 | print("Making graph") 123 | self.getSuperpixel() # get superpixel 124 | self.getCueValue_mean() # calculate average value of superpixels 125 | self.getConfidentMap() # meanwhile we get background seeds in this func 126 | 127 | if len(self.background_superseeds) == 0 or len(self.foreground_superseeds) == 0: 128 | print("Please enter at least one foreground and background seed.") 129 | return 130 | 131 | self.computeSigma() # get sigma 132 | 133 | # clear results 134 | if not os.path.exists(os.path.join('Results', self.ImageName)): 135 | os.mkdir(os.path.join('Results', self.ImageName)) 136 | for item in os.listdir(os.path.join('Results', self.ImageName)): 137 | itemsrc = os.path.join(os.path.join('Results', self.ImageName), item) 138 | os.remove(itemsrc) 139 | 140 | self.swap(10) # alpha-beta swap 141 | endtime = datetime.datetime.now() 142 | print("total run time: " + str((endtime - starttime).seconds)) 143 | 144 | self.getImage() 145 | 146 | def getSuperpixel(self): 147 | starttime = datetime.datetime.now() 148 | self.superpixel_segment = self.get_superpixel() 149 | # init_superpixel 150 | self.n_seg = np.amax(self.superpixel_segment) + 1 # 1 + num for superpixels 151 | self.num_seg = np.zeros(self.n_seg, dtype=int) # count for every cluster 152 | for x in range(0, len(self.superpixel_segment)): 153 | for y in range(0, len(self.superpixel_segment[0])): 154 | self.num_seg[self.superpixel_segment[x, y]] += 1 155 | self.super_label = np.ones(self.n_seg, dtype=int) * 0 156 | 157 | self.get_SuperEdge(self.superpixel_segment) 158 | 159 | endtime = datetime.datetime.now() 160 | print("get superpixel: " + str((endtime - starttime).seconds)) 161 | 162 | def getCueValue_mean(self): 163 | starttime = datetime.datetime.now() 164 | 165 | self.normal_map = self.normalMap() 166 | self.LAB_map = self.LABMap() 167 | self.depth_map = self.depth 168 | 169 | # normalMap for superpixel 170 | if self.NORMAL: 171 | self.ave_normal = np.zeros((self.n_seg, 3), dtype=float) 172 | for x in range(0, len(self.superpixel_segment)): 173 | for y in range(0, len(self.superpixel_segment[0])): 174 | self.ave_normal[self.superpixel_segment[x, y]] += self.normal_map[x, y] 175 | for i in range(0, self.n_seg): 176 | self.ave_normal[i] = self.normalizeVector(self.ave_normal[i]) 177 | 178 | # depthMap for superpixel 179 | self.ave_depth = np.zeros(self.n_seg, dtype=float) 180 | for x in range(0, len(self.superpixel_segment)): 181 | for y in range(0, len(self.superpixel_segment[0])): 182 | self.ave_depth[self.superpixel_segment[x, y]] += self.depth_map[x, y] 183 | for i in range(0, self.n_seg): 184 | self.ave_depth[i] /= self.num_seg[i] 185 | 186 | # LABMap for superpixel 187 | self.ave_LAB = np.zeros((self.n_seg, 3), dtype=float) 188 | for x in range(0, len(self.superpixel_segment)): 189 | for y in range(0, len(self.superpixel_segment[0])): 190 | self.ave_LAB[self.superpixel_segment[x, y]] += self.image[x, y] 191 | for i in range(0, self.n_seg): 192 | self.ave_LAB[i] /= self.num_seg[i] 193 | 194 | endtime = datetime.datetime.now() 195 | print("get mean-super-value for each cue: " + str((endtime - starttime).seconds)) 196 | 197 | def getConfidentMap(self): 198 | starttime = datetime.datetime.now() 199 | for coordinate in self.foreground_seeds: 200 | if self.superpixel_segment[coordinate[1] - 1, coordinate[0] - 1] not in self.foreground_superseeds: 201 | self.foreground_superseeds.append(self.superpixel_segment[coordinate[1] - 1, coordinate[0] - 1]) 202 | for coordinate in self.background_seeds: 203 | if self.superpixel_segment[coordinate[1] - 1, coordinate[0] - 1] not in self.background_superseeds: 204 | self.background_superseeds.append(self.superpixel_segment[coordinate[1] - 1, coordinate[0] - 1]) 205 | 206 | cost_func = lambda u, v, e, prev_e: e['cost'] 207 | 208 | #######For LAB foreground####### 209 | G_LAB = Graph() 210 | # add edges 211 | weight = [[] for i in range(0, self.n_seg)] 212 | aveMinWeight = 0 213 | for u in range(0, self.n_seg): 214 | minWeight = self.MAXIMUM 215 | for v in self.super_edge[u]: 216 | w = self.eu_dis(self.ave_LAB[u], self.ave_LAB[v]) 217 | weight[u].append((v, w)) 218 | if minWeight > w: 219 | minWeight = w 220 | aveMinWeight += minWeight 221 | aveMinWeight /= self.n_seg 222 | for u in range(0, self.n_seg): 223 | for v, w in weight[u]: 224 | if w < aveMinWeight: 225 | G_LAB.add_edge(u, v, {'cost': w / 3}) 226 | else: 227 | G_LAB.add_edge(u, v, {'cost': w}) 228 | 229 | for v in self.foreground_superseeds: 230 | G_LAB.add_edge(v, 's', {'cost': 0}) 231 | 232 | Lab_disFore = np.zeros(self.n_seg, dtype=float) 233 | Lab_disFore.fill(self.MAXIMUM) 234 | 235 | for v in range(0, self.n_seg): 236 | info = find_path(G_LAB, v, 's', cost_func=cost_func) 237 | Lab_disFore[v] = info.total_cost 238 | 239 | endtime = datetime.datetime.now() 240 | print("compute LAB disfore: " + str((endtime - starttime).seconds)) 241 | #######For Normal foreground####### 242 | if self.NORMAL: 243 | starttime = datetime.datetime.now() 244 | 245 | G_normal = Graph() 246 | # add edges 247 | weight = [[] for i in range(0, self.n_seg)] 248 | aveMinWeight = 0 249 | for u in range(0, self.n_seg): 250 | minWeight = self.MAXIMUM 251 | for v in self.super_edge[u]: 252 | w = 1 - self.cosine_sim(self.ave_normal[u], self.ave_normal[v]) 253 | weight[u].append((v, w)) 254 | if minWeight > w: 255 | minWeight = w 256 | aveMinWeight += minWeight 257 | aveMinWeight /= self.n_seg 258 | for u in range(0, self.n_seg): 259 | for v, w in weight[u]: 260 | if w < aveMinWeight: 261 | G_normal.add_edge(u, v, {'cost': w / 3}) 262 | else: 263 | G_normal.add_edge(u, v, {'cost': w}) 264 | 265 | for v in self.foreground_superseeds: 266 | G_normal.add_edge(v, 's', {'cost': 0}) 267 | 268 | Normal_disFore = np.zeros(self.n_seg, dtype=float) 269 | Normal_disFore.fill(self.MAXIMUM) 270 | for v in range(0, self.n_seg): 271 | info = find_path(G_normal, v, 's', cost_func=cost_func) 272 | Normal_disFore[v] = info.total_cost 273 | 274 | endtime = datetime.datetime.now() 275 | print("compute Normal disfore: " + str((endtime - starttime).seconds)) 276 | #######For depth foreground####### 277 | starttime = datetime.datetime.now() 278 | 279 | G_depth = Graph() 280 | # add edges 281 | weight = [[] for i in range(0, self.n_seg)] 282 | aveMinWeight = 0 283 | for u in range(0, self.n_seg): 284 | minWeight = self.MAXIMUM 285 | for v in self.super_edge[u]: 286 | w = np.abs(self.ave_depth[u] - self.ave_depth[v]) 287 | weight[u].append((v, w)) 288 | if minWeight > w: 289 | minWeight = w 290 | aveMinWeight += minWeight 291 | aveMinWeight /= self.n_seg 292 | 293 | for u in range(0, self.n_seg): 294 | for v, w in weight[u]: 295 | if w < aveMinWeight: 296 | G_depth.add_edge(u, v, {'cost': w / 3}) 297 | else: 298 | G_depth.add_edge(u, v, {'cost': w}) 299 | 300 | for v in self.foreground_superseeds: 301 | G_depth.add_edge(v, 's', {'cost': 0}) 302 | 303 | Depth_disFore = np.zeros(self.n_seg, dtype=float) 304 | Depth_disFore.fill(self.MAXIMUM) 305 | for v in range(0, self.n_seg): 306 | info = find_path(G_depth, v, 's', cost_func=cost_func) 307 | Depth_disFore[v] = info.total_cost 308 | 309 | Maxdis = 0 310 | for i in range(0, self.n_seg): 311 | if Maxdis < Depth_disFore[i]: 312 | Maxdis = Depth_disFore[i] 313 | 314 | endtime = datetime.datetime.now() 315 | print("compute Depth disfore: " + str((endtime - starttime).seconds)) 316 | #######get background_seed####### 317 | starttime = datetime.datetime.now() 318 | 319 | boundary_superpixel = [] 320 | for x in range(0, len(self.image)): 321 | if self.superpixel_segment[x, 0] not in boundary_superpixel: 322 | boundary_superpixel.append(self.superpixel_segment[x, 0]) 323 | if self.superpixel_segment[x, len(self.image[0]) - 1] not in boundary_superpixel: 324 | boundary_superpixel.append(self.superpixel_segment[x, len(self.image[0]) - 1]) 325 | for y in range(0, len(self.image[0])): 326 | if self.superpixel_segment[0, y] not in boundary_superpixel: 327 | boundary_superpixel.append(self.superpixel_segment[0, y]) 328 | if self.superpixel_segment[len(self.image) - 1, y] not in boundary_superpixel: 329 | boundary_superpixel.append(self.superpixel_segment[len(self.image) - 1, y]) 330 | for boundary in boundary_superpixel: 331 | if Depth_disFore[boundary] / Maxdis > 0.1 and boundary not in self.background_superseeds: 332 | self.background_superseeds.append(boundary) 333 | 334 | for v in self.background_superseeds: 335 | if self.NORMAL: 336 | G_normal.add_edge(v, 't', {'cost': 0}) 337 | G_LAB.add_edge(v, 't', {'cost': 0}) 338 | G_depth.add_edge(v, 't', {'cost': 0}) 339 | endtime = datetime.datetime.now() 340 | print("get background superseeds: " + str((endtime - starttime).seconds)) 341 | 342 | #######For LAB background####### 343 | starttime = datetime.datetime.now() 344 | Lab_disBack = np.zeros(self.n_seg, dtype=float) 345 | Lab_disBack.fill(self.MAXIMUM) 346 | for v in range(0, self.n_seg): 347 | info = find_path(G_LAB, v, 't', cost_func=cost_func) 348 | Lab_disBack[v] = info.total_cost 349 | 350 | endtime = datetime.datetime.now() 351 | print("compute LAB disback: " + str((endtime - starttime).seconds)) 352 | #######For Normal background####### 353 | if self.NORMAL: 354 | starttime = datetime.datetime.now() 355 | 356 | Normal_disBack = np.zeros(self.n_seg, dtype=float) 357 | Normal_disBack.fill(self.MAXIMUM) 358 | for v in range(0, self.n_seg): 359 | info = find_path(G_normal, v, 't', cost_func=cost_func) 360 | Normal_disBack[v] = info.total_cost 361 | 362 | endtime = datetime.datetime.now() 363 | print("compute Normal disback: " + str((endtime - starttime).seconds)) 364 | #######For Depth background####### 365 | starttime = datetime.datetime.now() 366 | 367 | Depth_disBack = np.zeros(self.n_seg, dtype=float) 368 | Depth_disBack.fill(self.MAXIMUM) 369 | for v in range(0, self.n_seg): 370 | info = find_path(G_depth, v, 't', cost_func=cost_func) 371 | Depth_disBack[v] = info.total_cost 372 | 373 | endtime = datetime.datetime.now() 374 | print("compute Depth disback: " + str((endtime - starttime).seconds)) 375 | #######confident map for each cue####### 376 | starttime = datetime.datetime.now() 377 | 378 | self.cfd_LAB = np.zeros(self.n_seg, dtype=float) 379 | if self.NORMAL: 380 | self.cfd_normal = np.zeros(self.n_seg, dtype=float) 381 | self.cfd_depth = np.zeros(self.n_seg, dtype=float) 382 | 383 | for i in range(0, self.n_seg): 384 | self.cfd_LAB[i] = Lab_disBack[i] / (Lab_disBack[i] + Lab_disFore[i]) 385 | if self.NORMAL: 386 | self.cfd_normal[i] = Normal_disBack[i] / (Normal_disBack[i] + Normal_disFore[i]) 387 | self.cfd_depth[i] = Depth_disBack[i] / (Depth_disBack[i] + Depth_disFore[i]) 388 | 389 | endtime = datetime.datetime.now() 390 | print("compute confident map: " + str((endtime - starttime).seconds)) 391 | 392 | def swap(self, MaxIteration=4): 393 | starttime = datetime.datetime.now() 394 | 395 | oldEnergy = self.computeEnerge() 396 | print("initial energe：" + str(oldEnergy)) 397 | 398 | for i in range(MaxIteration): 399 | 400 | self.oneSwapIteration(i) 401 | newEnergy = self.computeEnerge() 402 | 403 | print("iter" + str(i + 1) + " Energe：" + str(newEnergy)) 404 | if newEnergy >= oldEnergy: 405 | break 406 | oldEnergy = newEnergy 407 | endtime = datetime.datetime.now() 408 | print("alpha-beta-swap: " + str((endtime - starttime).seconds)) 409 | 410 | def oneSwapIteration(self, iteration): 411 | # self.get_seg_now(str(iteration) + "-00") 412 | old = self.computeEnerge() 413 | for i in range(0, 6): 414 | if not self.NORMAL and i >= 4: 415 | continue 416 | for j in range(i + 1, 6): 417 | if not self.NORMAL and j >= 4: 418 | continue 419 | self.alpha_beta_swap(i, j) 420 | new = self.computeEnerge() 421 | # if old > new: 422 | # self.get_seg_now(str(iteration) + "-" + str(i) + str(j)) 423 | old = new 424 | 425 | def alpha_beta_swap(self, alpha, beta): 426 | ###add nodes and edges for graphCuts### 427 | self.nodes = [] 428 | self.edges = [] 429 | reflect = [] ##discontinuity to continuity 430 | 431 | if alpha == 0: 432 | cap_source = 1 - self.cfd_LAB 433 | elif alpha == 1: 434 | cap_source = self.cfd_LAB 435 | elif alpha == 2: 436 | cap_source = 1 - self.cfd_depth 437 | elif alpha == 3: 438 | cap_source = self.cfd_depth 439 | elif alpha == 4: 440 | cap_source = 1 - self.cfd_normal 441 | elif alpha == 5: 442 | cap_source = self.cfd_normal 443 | 444 | if beta == 0: 445 | cap_sink = 1 - self.cfd_LAB 446 | elif beta == 1: 447 | cap_sink = self.cfd_LAB 448 | elif beta == 2: 449 | cap_sink = 1 - self.cfd_depth 450 | elif beta == 3: 451 | cap_sink = self.cfd_depth 452 | elif beta == 4: 453 | cap_sink = 1 - self.cfd_normal 454 | elif beta == 5: 455 | cap_sink = self.cfd_normal 456 | 457 | for i in range(self.n_seg): 458 | if self.super_label[i] == alpha or self.super_label[i] == beta: 459 | reflect.append(i) 460 | source_add = 0 461 | sink_add = 0 462 | for neighbor in self.super_edge[i]: 463 | if self.super_label[neighbor] != alpha and self.super_label[neighbor] != beta: 464 | source_add += self.smoothWeight(i, neighbor, alpha) 465 | sink_add += self.smoothWeight(i, neighbor, beta) 466 | self.nodes.append((reflect.index(i), cap_source[i] + source_add, cap_sink[i] + sink_add)) 467 | 468 | for n in self.nodes: 469 | u = reflect[n[0]] 470 | for v in self.super_edge[u]: 471 | if (self.super_label[v] == alpha or self.super_label[v] == beta) and v > u: 472 | # print(str(u) + "," + str(v)) 473 | weight = self.smoothWeight(u, v) 474 | self.edges.append((reflect.index(u), reflect.index(v), weight)) 475 | 476 | ####GraphCuts#### 477 | g = maxflow.Graph[float](len(self.nodes), len(self.edges)) 478 | 479 | nodelist = g.add_nodes(len(self.nodes)) 480 | for node in self.nodes: 481 | g.add_tedge(nodelist[node[0]], node[1], node[2]) 482 | 483 | for edge in self.edges: 484 | g.add_edge(edge[0], edge[1], edge[2], edge[2]) 485 | 486 | flow = g.maxflow() 487 | 488 | for vect in self.nodes: 489 | v = vect[0] 490 | if g.get_segment(v) == 0: # beta 491 | self.super_label[reflect[v]] = beta 492 | else: # alpha 493 | self.super_label[reflect[v]] = alpha 494 | 495 | def computeEnerge(self): 496 | return (self.giveDataEnerge() + self.giveSmoothEnerge()) 497 | 498 | def giveDataEnerge(self): 499 | energe = 0 500 | for i in range(self.n_seg): 501 | if self.super_label[i] == 0: 502 | energe += (1 - self.cfd_LAB[i]) 503 | elif self.super_label[i] == 1: 504 | energe += self.cfd_LAB[i] 505 | elif self.super_label[i] == 2: 506 | energe += (1 - self.cfd_depth[i]) 507 | elif self.super_label[i] == 3: 508 | energe += self.cfd_depth[i] 509 | elif self.super_label[i] == 4: 510 | energe += (1 - self.cfd_normal[i]) 511 | elif self.super_label[i] == 5: 512 | energe += self.cfd_normal[i] 513 | if energe != energe: 514 | print(str(i) + ":" + str(self.super_label[i])) 515 | break 516 | return energe 517 | 518 | def giveSmoothEnerge(self): # compute SmoothEnerge 519 | energe = 0 520 | for u in range(self.n_seg): 521 | for v in self.super_edge[u]: 522 | if v < u: 523 | continue 524 | if self.super_label[u] == self.super_label[v]: 525 | continue 526 | energe += self.smoothWeight(u, v) 527 | return energe 528 | 529 | def smoothWeight(self, u, v, alpha=-1): 530 | if u not in self.super_edge[v] or v not in self.super_edge[u]: 531 | return 0 532 | 533 | if alpha == -1: 534 | alpha = self.super_label[u] 535 | 536 | if alpha == 0 or alpha == 1: 537 | weight1 = self.lamda * np.e ** (-(self.eu_dis(self.ave_LAB[u], self.ave_LAB[v]) ** 2) / self.sigma1) 538 | elif alpha == 2 or alpha == 3: 539 | weight1 = self.lamda * np.e ** (-(np.abs(self.ave_depth[u] - self.ave_depth[v]) ** 2) / self.sigma2) 540 | elif alpha == 4 or alpha == 5: 541 | weight1 = self.lamda * np.e ** ( 542 | -((1 - self.cosine_sim(self.ave_normal[v], self.ave_normal[u])) ** 2) / self.sigma3) 543 | 544 | if self.super_label[v] == 0 or self.super_label[v] == 1: 545 | weight2 = self.lamda * np.e ** (-(self.eu_dis(self.ave_LAB[u], self.ave_LAB[v]) ** 2) / self.sigma1) 546 | elif self.super_label[v] == 2 or self.super_label[v] == 3: 547 | weight2 = self.lamda * np.e ** (-(np.abs(self.ave_depth[u] - self.ave_depth[v]) ** 2) / self.sigma2) 548 | elif self.super_label[v] == 4 or self.super_label[v] == 5: 549 | weight2 = self.lamda * np.e ** ( 550 | -((1 - self.cosine_sim(self.ave_normal[v], self.ave_normal[u])) ** 2) / self.sigma3) 551 | 552 | return min(weight1, weight2) 553 | 554 | @staticmethod 555 | def eu_dis(v1, v2): 556 | return np.sqrt((v1[0] - v2[0]) ** 2 + (v1[1] - v2[1]) ** 2 + (v1[2] - v2[2]) ** 2) 557 | 558 | @staticmethod 559 | def cosine_sim(vector1, vector2): 560 | return np.dot(vector1, vector2) / (np.linalg.norm(vector1) * (np.linalg.norm(vector2))) 561 | 562 | # SLIC SuperPixel 563 | def get_superpixel(self): 564 | segments = slic(self.image, n_segments=800) 565 | 566 | self.superpixel_image = img_as_ubyte(mark_boundaries(self.image, segments)) 567 | 568 | return segments 569 | 570 | def get_SuperEdge(self, segments): 571 | self.super_edge = [[] for _ in range(0, self.n_seg)] 572 | for x in range(0, len(segments)): 573 | for y in range(0, len(segments[0])): 574 | for i in range(-1, 2): 575 | for j in range(-1, 2): 576 | if i == 0 and j == 0: 577 | continue 578 | if (x + i < 0 or x + i >= len(segments) or y + j < 0 or y + j >= len( 579 | segments[0])): 580 | continue 581 | if (segments[x, y] != segments[x + i, y + j]): 582 | if segments[x, y] not in self.super_edge[ 583 | segments[x + i, y + j]]: 584 | self.super_edge[segments[x + i, y + j]].append( 585 | segments[x, y]) 586 | self.super_edge[segments[x, y]].append( 587 | segments[x + i, y + j]) 588 | 589 | # construct normalMap 590 | def normalMap(self): 591 | normal_map = np.zeros_like(self.image, dtype=float) 592 | width = self.depth.shape[1] 593 | height = self.depth.shape[0] 594 | 595 | for x in range(1, height - 1): 596 | for y in range(1, width - 1): 597 | dzdx = ((int)(self.depth[x + 1, y]) - (int)(self.depth[x - 1, y])) / 2 598 | dzdy = ((int)(self.depth[x, y + 1]) - (int)(self.depth[x, y - 1])) / 2 599 | d = (-dzdx, -dzdy, 1.0) 600 | n1 = self.normalizeVector(d) 601 | 602 | dzdxy = ((int)(self.depth[x + 1, y + 1]) - (int)(self.depth[x - 1, y - 1])) / 2.828 603 | dzdyx = ((int)(self.depth[x - 1, y + 1]) - (int)(self.depth[x + 1, y - 1])) / 2.828 604 | d = ((dzdyx - dzdxy), (-dzdxy - dzdyx), 2) 605 | n2 = self.normalizeVector(d) 606 | 607 | d = n1 + n2 608 | n = self.normalizeVector(d) 609 | 610 | normal_map[x, y] = n 611 | return normal_map 612 | 613 | # construct LABMap 614 | def LABMap(self): 615 | LAB_map_raw = cv2.cvtColor(self.image, cv2.COLOR_RGB2Lab) 616 | LAB_map = np.zeros_like(LAB_map_raw, dtype=np.int8) 617 | for i in range(len(LAB_map)): 618 | for j in range(len(LAB_map[0])): 619 | LAB_map[i, j][0] = LAB_map_raw[i, j][0] / 255 * 100 620 | LAB_map[i, j][1] = LAB_map_raw[i, j][1] - 128 621 | LAB_map[i, j][2] = LAB_map_raw[i, j][2] - 128 622 | return LAB_map 623 | 624 | @staticmethod 625 | def normalizeVector(d): 626 | return d / (np.sqrt(d[0] ** 2 + d[1] ** 2 + d[2] ** 2)) 627 | 628 | def compute_alpha_beta_energy(self, alpha, beta, reflect): 629 | energy = 0 630 | 631 | if alpha == 0: 632 | cap_source = 1 - self.cfd_LAB 633 | elif alpha == 1: 634 | cap_source = self.cfd_LAB 635 | elif alpha == 2: 636 | cap_source = 1 - self.cfd_depth 637 | elif alpha == 3: 638 | cap_source = self.cfd_depth 639 | elif alpha == 4: 640 | cap_source = 1 - self.cfd_normal 641 | elif alpha == 5: 642 | cap_source = self.cfd_normal 643 | 644 | if beta == 0: 645 | cap_sink = 1 - self.cfd_LAB 646 | elif beta == 1: 647 | cap_sink = self.cfd_LAB 648 | elif beta == 2: 649 | cap_sink = 1 - self.cfd_depth 650 | elif beta == 3: 651 | cap_sink = self.cfd_depth 652 | elif beta == 4: 653 | cap_sink = 1 - self.cfd_normal 654 | elif beta == 5: 655 | cap_sink = self.cfd_normal 656 | 657 | energyout = 0 658 | energyadd = 0 659 | energysmoothout = 0 660 | # for node in self.nodes: 661 | for i in range(self.n_seg): 662 | if self.super_label[i] == alpha: 663 | energy += cap_source[i] 664 | elif self.super_label[i] == beta: 665 | energy += cap_sink[i] 666 | else: 667 | if self.super_label[i] == 0: 668 | energyout += (1 - self.cfd_LAB[i]) 669 | elif self.super_label[i] == 1: 670 | energyout += self.cfd_LAB[i] 671 | elif self.super_label[i] == 2: 672 | energyout += (1 - self.cfd_depth[i]) 673 | elif self.super_label[i] == 3: 674 | energyout += self.cfd_depth[i] 675 | elif self.super_label[i] == 4: 676 | energyout += (1 - self.cfd_normal[i]) 677 | elif self.super_label[i] == 5: 678 | energyout += self.cfd_normal[i] 679 | 680 | if self.super_label[i] == alpha or self.super_label[i] == beta: 681 | for neighbor in self.super_edge[i]: 682 | if self.super_label[neighbor] != alpha and self.super_label[neighbor] != beta: 683 | energyadd += self.smoothWeight(i, neighbor) 684 | else: 685 | for neighbor in self.super_edge[i]: 686 | if self.super_label[neighbor] != alpha and self.super_label[neighbor] != beta and neighbor > i: 687 | energysmoothout += self.smoothWeight(i, neighbor) 688 | 689 | 690 | energysmooth = 0 691 | for edge in self.edges: 692 | if self.super_label[reflect[edge[0]]] != self.super_label[reflect[edge[1]]]: 693 | energysmooth += edge[2] 694 | energy += (energysmooth) 695 | return energy 696 | 697 | def getImage(self): 698 | #########check edge on superpixels######### 699 | temp_img = self.superpixel_image.copy() 700 | ave_cor = [[0, 0] for i in range(0, self.n_seg)] 701 | for x in range(0, len(self.superpixel_segment)): 702 | for y in range(0, len(self.superpixel_segment[0])): 703 | ave_cor[self.superpixel_segment[x, y]][0] += x 704 | ave_cor[self.superpixel_segment[x, y]][1] += y 705 | for i in range(0, self.n_seg): 706 | ave_cor[i] /= self.num_seg[i] 707 | for i in range(0, self.n_seg): 708 | for j in self.super_edge[i]: 709 | if j < i: 710 | continue 711 | cv2.line(temp_img, (int(ave_cor[i][1]), int(ave_cor[i][0])), (int(ave_cor[j][1]), int(ave_cor[j][0])), 712 | (0, 0, 255), 1) 713 | temp_img = temp_img.astype('uint8') 714 | cv2.imwrite(os.path.join('Results', self.ImageName, "edgeImg.jpg"), temp_img) 715 | ###get overlay### 716 | for x in range(0, len(self.superpixel_segment)): 717 | for y in range(0, len(self.superpixel_segment[0])): 718 | vect = self.superpixel_segment[x, y] 719 | if self.super_label[vect] % 2 == 0: 720 | self.segment_overlay[x, y] = (255, 255, 255) 721 | else: 722 | self.segment_overlay[x, y] = (0, 0, 0) 723 | 724 | ###get image with annotation### 725 | temp_img = self.image.copy() 726 | for x in range(0, len(temp_img)): 727 | for y in range(0, len(temp_img[0])): 728 | if self.superpixel_segment[x, y] in self.background_superseeds: 729 | temp_img[x, y] = (0, 0, 255) 730 | if self.superpixel_segment[x, y] in self.foreground_superseeds: 731 | temp_img[x, y] = (0, 255, 0) 732 | cv2.imwrite(os.path.join('Results', self.ImageName, "backseeds.jpg"), temp_img) 733 | 734 | ###get superpixels image### 735 | temp_img = self.superpixel_image.copy() 736 | cv2.imwrite(os.path.join('Results', self.ImageName, "superpixel.jpg"), temp_img) 737 | 738 | ###get LAB image### 739 | temp_img = self.LAB_map.copy() 740 | cv2.imwrite(os.path.join('Results', self.ImageName, "LABmap.jpg"), temp_img) 741 | 742 | ###get depth image### 743 | temp_img = self.depth_map / np.amax(self.depth_map) * 255 744 | cv2.imwrite(os.path.join('Results', self.ImageName, "depthmap.jpg"), temp_img) 745 | 746 | ###get normal image### 747 | if self.NORMAL: 748 | temp_img = self.normal_map * 255 749 | cv2.imwrite(os.path.join('Results', self.ImageName, "normalmap.jpg"), temp_img) 750 | 751 | ###get configent map of lab### 752 | temp_img = np.zeros_like(self.image) 753 | for x in range(0, len(self.superpixel_segment)): 754 | for y in range(0, len(self.superpixel_segment[0])): 755 | vect = self.superpixel_segment[x, y] 756 | temp_img[x, y] = self.cfd_LAB[vect] * 255 757 | cv2.imwrite(os.path.join('Results', self.ImageName, "confident_LAB.jpg"), temp_img) 758 | 759 | ###get confident map of depth### 760 | temp_img = np.zeros_like(self.image) 761 | for x in range(0, len(self.superpixel_segment)): 762 | for y in range(0, len(self.superpixel_segment[0])): 763 | vect = self.superpixel_segment[x, y] 764 | temp_img[x, y] = self.cfd_depth[vect] * 255 765 | cv2.imwrite(os.path.join('Results', self.ImageName, "confident_depth.jpg"), temp_img) 766 | ###get confident map of normal### 767 | if self.NORMAL: 768 | temp_img = np.zeros_like(self.image) 769 | for x in range(0, len(self.superpixel_segment)): 770 | for y in range(0, len(self.superpixel_segment[0])): 771 | vect = self.superpixel_segment[x, y] 772 | temp_img[x, y] = self.cfd_normal[vect] * 255 773 | cv2.imwrite(os.path.join('Results', self.ImageName, "confident_normal.jpg"), temp_img) 774 | 775 | ###get Results### 776 | temp_img = cv2.addWeighted(self.image, 0.2, self.segment_overlay, 0.8, 0.1) 777 | cv2.imwrite(os.path.join('Results', self.ImageName, "segment.jpg"), temp_img) 778 | ###Results with cues### 779 | self.segment_overlay = np.zeros_like(self.image) 780 | for x in range(0, len(self.superpixel_segment)): 781 | for y in range(0, len(self.superpixel_segment[0])): 782 | vect = self.superpixel_segment[x, y] 783 | if self.super_label[vect] == 0: 784 | self.segment_overlay[x, y] = (0, 0, 255) 785 | elif self.super_label[vect] == 1: 786 | self.segment_overlay[x, y] = (0, 255, 0) 787 | elif self.super_label[vect] == 2: 788 | self.segment_overlay[x, y] = (0, 255, 255) 789 | elif self.super_label[vect] == 3: 790 | self.segment_overlay[x, y] = (255, 0, 0) 791 | elif self.super_label[vect] == 4: 792 | self.segment_overlay[x, y] = (0, 153, 255) 793 | elif self.super_label[vect] == 5: 794 | self.segment_overlay[x, y] = (255, 0, 153) 795 | 796 | temp_img = cv2.addWeighted(self.image, 0.2, self.segment_overlay, 0.8, 0.1) 797 | cv2.imwrite(os.path.join('Results', self.ImageName, "segment_cues.jpg"), temp_img) 798 | 799 | temp_img = np.zeros_like(self.image) 800 | w = np.amax(self.depth_map) 801 | for x in range(0, len(self.superpixel_segment)): 802 | for y in range(0, len(self.superpixel_segment[0])): 803 | temp_img[x, y] = self.ave_depth[self.superpixel_segment[x, y]] / w * 255 804 | cv2.imwrite(os.path.join('Results', self.ImageName, "ave_depth.jpg"), temp_img) 805 | 806 | def get_seg_now(self, path): # get current seg 807 | for x in range(0, len(self.superpixel_segment)): 808 | for y in range(0, len(self.superpixel_segment[0])): 809 | vect = self.superpixel_segment[x, y] 810 | if self.super_label[vect] == 0: 811 | self.segment_overlay[x, y] = (0, 0, 255) 812 | elif self.super_label[vect] == 1: 813 | self.segment_overlay[x, y] = (0, 255, 0) 814 | elif self.super_label[vect] == 2: 815 | self.segment_overlay[x, y] = (0, 255, 255) 816 | elif self.super_label[vect] == 3: 817 | self.segment_overlay[x, y] = (255, 0, 0) 818 | elif self.super_label[vect] == 4: 819 | self.segment_overlay[x, y] = (0, 153, 255) 820 | elif self.super_label[vect] == 5: 821 | self.segment_overlay[x, y] = (255, 0, 153) 822 | temp = img_as_ubyte(mark_boundaries(self.get_image_with_overlay(self.segmented), self.superpixel_segment)) 823 | cv2.imwrite(os.path.join('Results', self.ImageName, path + ".jpg"), temp) 824 | self.segment_overlay = np.zeros_like(self.segment_overlay) 825 | -------------------------------------------------------------------------------- /Imgs/Depth/8_08-40-06_Depth.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZVsion/rgbd_image_segmentation/96b8c5dfe4b22231b8eb13e5fc34d3c0c3231e26/Imgs/Depth/8_08-40-06_Depth.png -------------------------------------------------------------------------------- /Imgs/RGB/8_08-40-06.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZVsion/rgbd_image_segmentation/96b8c5dfe4b22231b8eb13e5fc34d3c0c3231e26/Imgs/RGB/8_08-40-06.jpg -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Interactive-Segmentation-on-RGBD-Images-via-Cue-Selection 2 | Python implementation of paper "Interactive Segmentation on RGBD Images via Cue Selection", CVPR 2016. 3 | 4 | ![top](doc/top.png) 5 | 6 | ### Dependencies 7 | ``` 8 | Python3.x (Tested with 3.5) 9 | opencv 10 | PyQt5 11 | scikit-learn 12 | PyMaxflow 13 | Dijkstar 14 | ``` 15 | 16 | ### Installation 17 | To use this code, please do: 18 | 19 | 0. Clone the repo: 20 | ```Shell 21 | git clone https://github.com/ZVsion/rgbd_image_segmentation 22 | cd rgbd_image_segmentation 23 | ``` 24 | 25 | 1. To try the demo, please run: 26 | ```Shell 27 | python main.py 28 | ``` 29 | 30 | If installed correctly, the GUI should look like this: 31 |

32 | 33 | You can add some foreground and background seeds in image, and then click ``Segment Image``. After a few seconds, 34 | the result should look like this: 35 |

36 | 37 | Each color's label is presented as follow: 38 | 39 | ![label](doc/label.png) 40 | 41 | 42 | For more details about this algorithm, please see in [paper](http://openaccess.thecvf.com/content_cvpr_2016/papers/Feng_Interactive_Segmentation_on_CVPR_2016_paper.pdf) and [code](https://github.com/ZVsion/rgbd_image_segmentation/blob/master/GraphMaker.py). 43 | 44 | Noted that we do not implement Grabcut postprocessing here. 45 | 46 | -------------------------------------------------------------------------------- /Results/.gitignore: -------------------------------------------------------------------------------- 1 | * 2 | !.gitignore 3 | -------------------------------------------------------------------------------- /doc/gui1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZVsion/rgbd_image_segmentation/96b8c5dfe4b22231b8eb13e5fc34d3c0c3231e26/doc/gui1.png -------------------------------------------------------------------------------- /doc/label.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZVsion/rgbd_image_segmentation/96b8c5dfe4b22231b8eb13e5fc34d3c0c3231e26/doc/label.png -------------------------------------------------------------------------------- /doc/segment_cues.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZVsion/rgbd_image_segmentation/96b8c5dfe4b22231b8eb13e5fc34d3c0c3231e26/doc/segment_cues.jpg -------------------------------------------------------------------------------- /doc/top.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZVsion/rgbd_image_segmentation/96b8c5dfe4b22231b8eb13e5fc34d3c0c3231e26/doc/top.png -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from PyQt5.QtGui import * 3 | from PyQt5.QtWidgets import * 4 | from PyQt5.QtCore import * 5 | import cv2 6 | import GraphMaker 7 | 8 | class Example(QWidget): 9 | 10 | def __init__(self): 11 | super().__init__() 12 | 13 | self.graph_maker = GraphMaker.GraphMaker() 14 | self.seed_type = 1 #annotation type 15 | 16 | self.initUI() 17 | 18 | def initUI(self): 19 | self.resize(1000, 600) 20 | 21 | clearButton = QPushButton("Clear All Seeds") 22 | clearButton.setStyleSheet("background-color:white") 23 | clearButton.clicked.connect(self.on_clear) 24 | 25 | segmentButton = QPushButton("Segment Image") 26 | segmentButton.setStyleSheet("background-color:white") 27 | segmentButton.clicked.connect(self.on_segment) 28 | 29 | StateLine = QLabel() 30 | StateLine.setText("Click or Drag Left Mouse Button for Foreground Annotation,and Right Mouse Button for Background.") 31 | palette = QPalette() 32 | palette.setColor(StateLine.foregroundRole(), Qt.red) 33 | StateLine.setPalette(palette) 34 | 35 | hbox = QHBoxLayout() 36 | hbox.addWidget(clearButton) 37 | hbox.addWidget(segmentButton) 38 | hbox.addStretch(1) 39 | hbox.addWidget(StateLine) 40 | hbox.addStretch(1) 41 | 42 | self.seedLabel = QLabel() 43 | self.seedLabel.setPixmap(QPixmap.fromImage( 44 | self.get_qimage(self.graph_maker.get_image_with_overlay(self.graph_maker.seeds)))) 45 | self.seedLabel.setAlignment(Qt.AlignCenter) 46 | self.seedLabel.mousePressEvent = self.mouse_down 47 | self.seedLabel.mouseMoveEvent = self.mouse_drag 48 | 49 | self.segmentLabel = QLabel() 50 | self.segmentLabel.setPixmap(QPixmap.fromImage( 51 | self.get_qimage(self.graph_maker.get_image_with_overlay(self.graph_maker.segmented)))) 52 | self.segmentLabel.setAlignment(Qt.AlignCenter) 53 | 54 | scroll1 = QScrollArea() 55 | scroll1.setWidgetResizable(False) 56 | scroll1.setWidget(self.seedLabel) 57 | scroll1.setAlignment(Qt.AlignCenter) 58 | 59 | scroll2 = QScrollArea() 60 | scroll2.setWidgetResizable(False) 61 | scroll2.setWidget(self.segmentLabel) 62 | scroll2.setAlignment(Qt.AlignCenter) 63 | 64 | imagebox = QHBoxLayout() 65 | imagebox.addWidget(scroll1) 66 | imagebox.addWidget(scroll2) 67 | 68 | vbox = QVBoxLayout() 69 | 70 | vbox.addLayout(hbox) 71 | vbox.addLayout(imagebox) 72 | 73 | self.setLayout(vbox) 74 | 75 | self.setWindowTitle('Buttons') 76 | self.show() 77 | @staticmethod 78 | def get_qimage(cvimage): 79 | height, width, bytes_per_pix = cvimage.shape 80 | bytes_per_line = width * bytes_per_pix 81 | cv2.cvtColor(cvimage, cv2.COLOR_BGR2RGB, cvimage) 82 | return QImage(cvimage.data, width, height, bytes_per_line, QImage.Format_RGB888) 83 | 84 | def mouse_down(self, event): 85 | if event.button() == Qt.LeftButton: 86 | self.seed_type = 1 87 | elif event.button() == Qt.RightButton: 88 | self.seed_type = 0 89 | print(str(event.x()) + "," + str(event.y())) 90 | self.graph_maker.add_seed(event.x(), event.y(), self.seed_type) 91 | self.seedLabel.setPixmap(QPixmap.fromImage( 92 | self.get_qimage(self.graph_maker.get_image_with_overlay(self.graph_maker.seeds)))) 93 | 94 | def mouse_drag(self, event): 95 | self.graph_maker.add_seed(event.x(), event.y(), self.seed_type) 96 | self.seedLabel.setPixmap(QPixmap.fromImage( 97 | self.get_qimage(self.graph_maker.get_image_with_overlay(self.graph_maker.seeds)))) 98 | 99 | @pyqtSlot() 100 | def on_segment(self): 101 | # caculate 102 | self.graph_maker.create_graph() 103 | 104 | self.segmentLabel.setPixmap(QPixmap.fromImage( 105 | self.get_qimage(self.graph_maker.get_image_with_overlay(self.graph_maker.segmented)))) 106 | 107 | @pyqtSlot() 108 | def on_clear(self): 109 | self.graph_maker.clear_seeds() 110 | self.seedLabel.setPixmap(QPixmap.fromImage( 111 | self.get_qimage(self.graph_maker.get_image_with_overlay(self.graph_maker.seeds)))) 112 | 113 | if __name__ == '__main__': 114 | app = QApplication(sys.argv) 115 | ex = Example() 116 | sys.exit(app.exec_()) --------------------------------------------------------------------------------