Photoshop中磁力套索的一种简陋实现(基于Python)

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

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\n*.jpg\n*.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')

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()