Photoshop中磁力套索的一種簡陋實現(xiàn)(基于Python)

輸入:
 s // 種子點
 l(q,r) // 計算局部cost

數(shù)據(jù)結(jié)構(gòu):
 L // 當(dāng)前待處理的像素
 N(q) // 當(dāng)前像素相鄰的像素
 e(q) // 標(biāo)記一個像素是否已經(jīng)做過相鄰像素展開的Bool函數(shù)
 g(q) // 從s到q的總cost

from __future__ import division
import cv2
import numpy as np

SQRT_0_5 = 0.70710678118654757

class Livewire():
 """
 A simple livewire implementation for verification using 
 1. Canny edge detector + gradient magnitude + gradient direction
 2. Dijkstra algorithm
 """
 
 def __init__(self, image):
 self.image = image
 self.x_lim = image.shape[0]
 self.y_lim = image.shape[1]
 # The values in cost matrix ranges from 0~1
 self.cost_edges = 1 - cv2.Canny(image, 85, 170)/255.0
 self.grad_x, self.grad_y, self.grad_mag = self._get_grad(image)
 self.cost_grad_mag = 1 - self.grad_mag/np.max(self.grad_mag)
 # Weight for (Canny edges, gradient magnitude, gradient direction)
 self.weight = (0.425, 0.425, 0.15)
 
 self.n_pixs = self.x_lim * self.y_lim
 self.n_processed = 0
 
 @classmethod
 def _get_grad(cls, image):
 """
 Return the gradient magnitude of the image using Sobel operator
 """
 rgb = True if len(image.shape) > 2 else False
 grad_x = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=3)
 grad_y = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=3)
 if rgb:
 # A very rough approximation for quick verification...
 grad_x = np.max(grad_x, axis=2)
 grad_y = np.max(grad_y, axis=2)
 
 grad_mag = np.sqrt(grad_x**2+grad_y**2)
 grad_x /= grad_mag
 grad_y /= grad_mag
 
 return grad_x, grad_y, grad_mag
 
 def _get_neighbors(self, p):
 """
 Return 8 neighbors around the pixel p
 """
 x, y = p
 x0 = 0 if x == 0 else x - 1
 x1 = self.x_lim if x == self.x_lim - 1 else x + 2
 y0 = 0 if y == 0 else y - 1
 y1 = self.y_lim if y == self.y_lim - 1 else y + 2
 
 return [(x, y) for x in xrange(x0, x1) for y in xrange(y0, y1) if (x, y) != p]
 
 def _get_grad_direction_cost(self, p, q):
 """
 Calculate the gradient changes refer to the link direction
 """
 dp = (self.grad_y[p[0]][p[1]], -self.grad_x[p[0]][p[1]])
 dq = (self.grad_y[q[0]][q[1]], -self.grad_x[q[0]][q[1]])
 
 l = np.array([q[0]-p[0], q[1]-p[1]], np.float)
 if 0 not in l:
 l *= SQRT_0_5
 
 dp_l = np.dot(dp, l)
 l_dq = np.dot(l, dq)
 if dp_l < 0:
 dp_l = -dp_l
 l_dq = -l_dq
 
 # 2/3pi * ...
 return 0.212206590789 * (np.arccos(dp_l)+np.arccos(l_dq))
 
 def _local_cost(self, p, q):
 """
 1. Calculate the Canny edges & gradient magnitude cost taking into account Euclidean distance
 2. Combine with gradient direction
 Assumption: p & q are neighbors
 """
 diagnol = q[0] == p[0] or q[1] == p[1]
 
 # c0, c1 and c2 are costs from Canny operator, gradient magnitude and gradient direction respectively
 if diagnol:
 c0 = self.cost_edges[p[0]][p[1]]-SQRT_0_5*(self.cost_edges[p[0]][p[1]]-self.cost_edges[q[0]][q[1]])
 c1 = self.cost_grad_mag[p[0]][p[1]]-SQRT_0_5*(self.cost_grad_mag[p[0]][p[1]]-self.cost_grad_mag[q[0]][q[1]])
 c2 = SQRT_0_5 * self._get_grad_direction_cost(p, q)
 else:
 c0 = self.cost_edges[q[0]][q[1]]
 c1 = self.cost_grad_mag[q[0]][q[1]]
 c2 = self._get_grad_direction_cost(p, q)
 
 if np.isnan(c2):
 c2 = 0.0
 
 w0, w1, w2 = self.weight
 cost_pq = w0*c0 + w1*c1 + w2*c2
 
 return cost_pq * cost_pq

 def get_path_matrix(self, seed):
 """
 Get the back tracking matrix of the whole image from the cost matrix
 """
 neighbors = [] # 8 neighbors of the pixel being processed
 processed = set() # Processed point
 cost = {seed: 0.0} # Accumulated cost, initialized with seed to itself
 paths = {}

 self.n_processed = 0
 
 while cost:
 # Expand the minimum cost point
 p = min(cost, key=cost.get)
 neighbors = self._get_neighbors(p)
 processed.add(p)

 # Record accumulated costs and back tracking point for newly expanded points
 for q in [x for x in neighbors if x not in processed]:
 temp_cost = cost[p] + self._local_cost(p, q)
 if q in cost:
 if temp_cost < cost[q]:
 cost.pop(q)
 else:
 cost[q] = temp_cost
 processed.add(q)
 paths[q] = p
 
 # Pop traversed points
 cost.pop(p)
 
 self.n_processed += 1
 
 return paths

livewire.py

 from __future__ import division
import time
import cv2
from PyQt4 import QtGui, QtCore
from threading import Thread
from livewire import Livewire

class ImageWin(QtGui.QWidget):
 def __init__(self):
 super(ImageWin, self).__init__()
 self.setupUi()
 self.active = False
 self.seed_enabled = True
 self.seed = None
 self.path_map = {}
 self.path = []
 
 def setupUi(self):
 self.hbox = QtGui.QVBoxLayout(self)
 
 # Load and initialize image
 self.image_path = ''
 while self.image_path == '':
 self.image_path = QtGui.QFileDialog.getOpenFileName(self, '', '', '(*.bmp *.jpg *.png)')
 self.image = QtGui.QPixmap(self.image_path)
 self.cv2_image = cv2.imread(str(self.image_path))
 self.lw = Livewire(self.cv2_image)
 self.w, self.h = self.image.width(), self.image.height()
 
 self.canvas = QtGui.QLabel(self)
 self.canvas.setMouseTracking(True)
 self.canvas.setPixmap(self.image)
 
 self.status_bar = QtGui.QStatusBar(self)
 self.status_bar.showMessage('Left click to set a seed')
 
 self.hbox.addWidget(self.canvas)
 self.hbox.addWidget(self.status_bar)
 self.setLayout(self.hbox)
 
 def mousePressEvent(self, event): 
 if self.seed_enabled:
 pos = event.pos()
 x, y = pos.x()-self.canvas.x(), pos.y()-self.canvas.y()
 
 if x < 0:
 x = 0
 if x >= self.w:
 x = self.w - 1
 if y < 0:
 y = 0
 if y >= self.h:
 y = self.h - 1

 # Get the mouse cursor position
 p = y, x
 seed = self.seed
 
 # Export bitmap
 if event.buttons() == QtCore.Qt.MidButton:
 filepath = QtGui.QFileDialog.getSaveFileName(self, 'Save image audio to', '', '*.bmp
*.jpg
*.png')
 image = self.image.copy()
 
 draw = QtGui.QPainter()
 draw.begin(image)
 draw.setPen(QtCore.Qt.blue)
 if self.path_map:
 while p != seed:
 draw.drawPoint(p[1], p[0])
 for q in self.lw._get_neighbors(p):
 draw.drawPoint(q[1], q[0])
 p = self.path_map[p]
 if self.path:
 draw.setPen(QtCore.Qt.green)
 for p in self.path:
 draw.drawPoint(p[1], p[0])
 for q in self.lw._get_neighbors(p):
 draw.drawPoint(q[1], q[0])
 draw.end()
 
 image.save(filepath, quality=100)
 
 else:
 self.seed = p
 
 if self.path_map:
 while p != seed:
 p = self.path_map[p]
 self.path.append(p)
 
 # Calculate path map
 if event.buttons() == QtCore.Qt.LeftButton:
 Thread(target=self._cal_path_matrix).start()
 Thread(target=self._update_path_map_progress).start()
 
 # Finish current task and reset
 elif event.buttons() == QtCore.Qt.RightButton:
 self.path_map = {}
 self.status_bar.showMessage('Left click to set a seed')
 self.active = False
 
 def mouseMoveEvent(self, event):
 if self.active and event.buttons() == QtCore.Qt.NoButton:
 pos = event.pos()
 x, y = pos.x()-self.canvas.x(), pos.y()-self.canvas.y()

 if x < 0 or x >= self.w or y < 0 or y >= self.h:
 pass
 else:
 # Draw livewire
 p = y, x
 path = []
 while p != self.seed:
 p = self.path_map[p]
 path.append(p)
 
 image = self.image.copy()
 draw = QtGui.QPainter()
 draw.begin(image)
 draw.setPen(QtCore.Qt.blue)
 for p in path:
 draw.drawPoint(p[1], p[0])
 if self.path:
 draw.setPen(QtCore.Qt.green)
 for p in self.path:
 draw.drawPoint(p[1], p[0])
 draw.end()
 self.canvas.setPixmap(image)
 
 def _cal_path_matrix(self):
 self.seed_enabled = False
 self.active = False
 self.status_bar.showMessage('Calculating path map...')
 path_matrix = self.lw.get_path_matrix(self.seed)
 self.status_bar.showMessage(r'Left: new seed / Right: finish')
 self.seed_enabled = True
 self.active = True
 
 self.path_map = path_matrix
 
 def _update_path_map_progress(self):
 while not self.seed_enabled:
 time.sleep(0.1)
 message = 'Calculating path map... {:.1f}%'.format(self.lw.n_processed/self.lw.n_pixs*100.0)
 self.status_bar.showMessage(message)
 self.status_bar.showMessage(r'Left: new seed / Right: finish')

gui.py

import sys
from PyQt4 import QtGui
from gui import ImageWin

def main():
 app = QtGui.QApplication(sys.argv)
 window = ImageWin()
 window.setMouseTracking(True)
 window.setWindowTitle('Livewire Demo')
 window.show()
 window.setFixedSize(window.size())
 sys.exit(app.exec_())

if __name__ == '__main__':
 main() 

main.py