247 lines
9.3 KiB
Python
247 lines
9.3 KiB
Python
import cv2
|
|
import numpy as np
|
|
import math
|
|
import Polygon as plg
|
|
|
|
|
|
def watershed1(image, viz=False):
|
|
boxes = []
|
|
if len(image.shape) == 3:
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
else:
|
|
gray = image
|
|
if viz:
|
|
cv2.imshow("gray", gray)
|
|
cv2.waitKey()
|
|
ret, binary = cv2.threshold(gray, 0.6 * np.max(gray), 255, cv2.THRESH_BINARY)
|
|
if viz:
|
|
cv2.imshow("binary", binary)
|
|
cv2.waitKey()
|
|
# 形态学操作，进一步消除图像中噪点
|
|
kernel = np.ones((3, 3), np.uint8)
|
|
# kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
|
|
mb = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel, iterations=2) # iterations连续两次开操作
|
|
sure_bg = cv2.dilate(mb, kernel, iterations=3) # 3次膨胀,可以获取到大部分都是背景的区域
|
|
if viz:
|
|
cv2.imshow("sure_bg", sure_bg)
|
|
cv2.waitKey()
|
|
|
|
# 距离变换
|
|
dist = cv2.distanceTransform(mb, cv2.DIST_L2, 5)
|
|
if viz:
|
|
cv2.imshow("dist", dist)
|
|
cv2.waitKey()
|
|
ret, sure_fg = cv2.threshold(dist, 0.2 * np.max(dist), 255, cv2.THRESH_BINARY)
|
|
surface_fg = np.uint8(sure_fg) # 保持色彩空间一致才能进行运算，现在是背景空间为整型空间，前景为浮点型空间，所以进行转换
|
|
if viz:
|
|
cv2.imshow("surface_fg", surface_fg)
|
|
cv2.waitKey()
|
|
unknown = cv2.subtract(sure_bg, surface_fg)
|
|
# 获取maskers,在markers中含有种子区域
|
|
ret, markers = cv2.connectedComponents(surface_fg)
|
|
|
|
# 分水岭变换
|
|
markers = markers + 1
|
|
markers[unknown == 255] = 0
|
|
|
|
if viz:
|
|
color_markers = np.uint8(markers)
|
|
color_markers = cv2.applyColorMap(color_markers, cv2.COLORMAP_JET)
|
|
cv2.imshow("color_markers", color_markers)
|
|
cv2.waitKey()
|
|
|
|
markers = cv2.watershed(image, markers=markers)
|
|
image[markers == -1] = [0, 0, 255]
|
|
if viz:
|
|
cv2.imshow("image", image)
|
|
cv2.waitKey()
|
|
for i in range(2, np.max(markers) + 1):
|
|
np_contours = np.roll(np.array(np.where(markers == i)), 1, axis=0).transpose().reshape(-1, 2)
|
|
# print(np_contours.shape)
|
|
rectangle = cv2.minAreaRect(np_contours)
|
|
box = cv2.boxPoints(rectangle)
|
|
w, h = np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[1] - box[2])
|
|
box_ratio = max(w, h) / (min(w, h) + 1e-5)
|
|
if abs(1 - box_ratio) <= 0.1:
|
|
l, r = min(np_contours[:, 0]), max(np_contours[:, 0])
|
|
t, b = min(np_contours[:, 1]), max(np_contours[:, 1])
|
|
box = np.array([[l, t], [r, t], [r, b], [l, b]], dtype=np.float32)
|
|
|
|
# make clock-wise order
|
|
startidx = box.sum(axis=1).argmin()
|
|
box = np.roll(box, 4 - startidx, 0)
|
|
box = np.array(box)
|
|
boxes.append(box)
|
|
return np.array(boxes)
|
|
|
|
|
|
def getDetCharBoxes_core(textmap, text_threshold=0.5, low_text=0.4):
|
|
# prepare data
|
|
textmap = textmap.copy()
|
|
img_h, img_w = textmap.shape
|
|
|
|
""" labeling method """
|
|
ret, text_score = cv2.threshold(textmap, low_text, 1, 0)
|
|
nLabels, labels, stats, centroids = cv2.connectedComponentsWithStats(text_score.astype(np.uint8),
|
|
connectivity=4)
|
|
|
|
det = []
|
|
mapper = []
|
|
for k in range(1, nLabels):
|
|
# size filtering
|
|
size = stats[k, cv2.CC_STAT_AREA]
|
|
if size < 10: continue
|
|
|
|
# thresholding
|
|
if np.max(textmap[labels == k]) < text_threshold: continue
|
|
|
|
# make segmentation map
|
|
segmap = np.zeros(textmap.shape, dtype=np.uint8)
|
|
segmap[labels == k] = 255
|
|
# segmap[np.logical_and(link_score == 1, text_score == 0)] = 0 # remove link area
|
|
x, y = stats[k, cv2.CC_STAT_LEFT], stats[k, cv2.CC_STAT_TOP]
|
|
w, h = stats[k, cv2.CC_STAT_WIDTH], stats[k, cv2.CC_STAT_HEIGHT]
|
|
niter = int(math.sqrt(size * min(w, h) / (w * h)) * 2)
|
|
sx, ex, sy, ey = x - niter, x + w + niter + 1, y - niter, y + h + niter + 1
|
|
# boundary check
|
|
if sx < 0: sx = 0
|
|
if sy < 0: sy = 0
|
|
if ex >= img_w: ex = img_w
|
|
if ey >= img_h: ey = img_h
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1 + niter, 1 + niter))
|
|
segmap[sy:ey, sx:ex] = cv2.dilate(segmap[sy:ey, sx:ex], kernel)
|
|
|
|
# make box
|
|
np_contours = np.roll(np.array(np.where(segmap != 0)), 1, axis=0).transpose().reshape(-1, 2)
|
|
rectangle = cv2.minAreaRect(np_contours)
|
|
box = cv2.boxPoints(rectangle)
|
|
|
|
# align diamond-shape
|
|
w, h = np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[1] - box[2])
|
|
box_ratio = max(w, h) / (min(w, h) + 1e-5)
|
|
if abs(1 - box_ratio) <= 0.1:
|
|
l, r = min(np_contours[:, 0]), max(np_contours[:, 0])
|
|
t, b = min(np_contours[:, 1]), max(np_contours[:, 1])
|
|
box = np.array([[l, t], [r, t], [r, b], [l, b]], dtype=np.float32)
|
|
|
|
# make clock-wise order
|
|
startidx = box.sum(axis=1).argmin()
|
|
box = np.roll(box, 4 - startidx, 0)
|
|
box = np.array(box)
|
|
|
|
det.append(box)
|
|
mapper.append(k)
|
|
|
|
return det, labels, mapper
|
|
|
|
|
|
def watershed2(image, viz=False):
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
gray = np.float32(gray) / 255.0
|
|
boxes, _, _ = getDetCharBoxes_core(gray)
|
|
return np.array(boxes)
|
|
|
|
|
|
def watershed(oriimage, image, viz=False):
|
|
# viz = True
|
|
boxes = []
|
|
if len(image.shape) == 3:
|
|
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
|
|
else:
|
|
gray = image
|
|
if viz:
|
|
cv2.imshow("gray", gray)
|
|
cv2.waitKey()
|
|
ret, binary = cv2.threshold(gray, 0.2 * np.max(gray), 255, cv2.THRESH_BINARY)
|
|
if viz:
|
|
cv2.imshow("binary", binary)
|
|
cv2.waitKey()
|
|
# 形态学操作，进一步消除图像中噪点
|
|
kernel = np.ones((3, 3), np.uint8)
|
|
# kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
|
|
mb = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel, iterations=2) # iterations连续两次开操作
|
|
sure_bg = cv2.dilate(mb, kernel, iterations=3) # 3次膨胀,可以获取到大部分都是背景的区域
|
|
sure_bg = mb
|
|
if viz:
|
|
cv2.imshow("sure_bg", mb)
|
|
cv2.waitKey()
|
|
# 距离变换
|
|
# dist = cv2.distanceTransform(mb, cv2.DIST_L2, 5)
|
|
# if viz:
|
|
# cv2.imshow("dist", dist)
|
|
# cv2.waitKey()
|
|
ret, sure_fg = cv2.threshold(gray, 0.6 * gray.max(), 255, cv2.THRESH_BINARY)
|
|
surface_fg = np.uint8(sure_fg) # 保持色彩空间一致才能进行运算，现在是背景空间为整型空间，前景为浮点型空间，所以进行转换
|
|
if viz:
|
|
cv2.imshow("surface_fg", surface_fg)
|
|
cv2.waitKey()
|
|
unknown = cv2.subtract(sure_bg, surface_fg)
|
|
# 获取maskers,在markers中含有种子区域
|
|
ret, markers = cv2.connectedComponents(surface_fg)
|
|
|
|
nLabels, labels, stats, centroids = cv2.connectedComponentsWithStats(surface_fg,
|
|
connectivity=4)
|
|
# 分水岭变换
|
|
markers = labels.copy() + 1
|
|
# markers = markers+1
|
|
markers[unknown == 255] = 0
|
|
|
|
if viz:
|
|
color_markers = np.uint8(markers)
|
|
color_markers = color_markers / (color_markers.max() / 255)
|
|
color_markers = np.uint8(color_markers)
|
|
color_markers = cv2.applyColorMap(color_markers, cv2.COLORMAP_JET)
|
|
cv2.imshow("color_markers", color_markers)
|
|
cv2.waitKey()
|
|
# a = cv2.applyColorMap(gray, cv2.COLORMAP_JET)
|
|
markers = cv2.watershed(oriimage, markers=markers)
|
|
oriimage[markers == -1] = [0, 0, 255]
|
|
|
|
if viz:
|
|
color_markers = np.uint8(markers + 1)
|
|
color_markers = color_markers / (color_markers.max() / 255)
|
|
color_markers = np.uint8(color_markers)
|
|
color_markers = cv2.applyColorMap(color_markers, cv2.COLORMAP_JET)
|
|
cv2.imshow("color_markers1", color_markers)
|
|
cv2.waitKey()
|
|
|
|
if viz:
|
|
cv2.imshow("image", oriimage)
|
|
cv2.waitKey()
|
|
for i in range(2, np.max(markers) + 1):
|
|
np_contours = np.roll(np.array(np.where(markers == i)), 1, axis=0).transpose().reshape(-1, 2)
|
|
# segmap = np.zeros(gray.shape, dtype=np.uint8)
|
|
# segmap[markers == i] = 255
|
|
# size = np_contours.shape[0]
|
|
# x, y, w, h = cv2.boundingRect(np_contours)
|
|
# if w == 0 or h == 0:
|
|
# continue
|
|
#
|
|
# niter = int(math.sqrt(size * min(w, h) / (w * h)) * 2)
|
|
# sx, ex, sy, ey = x - niter, x + w + niter + 1, y - niter, y + h + niter + 1
|
|
# # boundary check
|
|
# if sx < 0: sx = 0
|
|
# if sy < 0: sy = 0
|
|
# if ex >= gray.shape[1]: ex = gray.shape[1]
|
|
# if ey >= gray.shape[0]: ey = gray.shape[0]
|
|
# kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1 + niter, 1 + niter))
|
|
# segmap[sy:ey, sx:ex] = cv2.dilate(segmap[sy:ey, sx:ex], kernel)
|
|
# np_contours = np.roll(np.array(np.where(segmap != 0)), 1, axis=0).transpose().reshape(-1, 2)
|
|
rectangle = cv2.minAreaRect(np_contours)
|
|
box = cv2.boxPoints(rectangle)
|
|
|
|
startidx = box.sum(axis=1).argmin()
|
|
box = np.roll(box, 4 - startidx, 0)
|
|
poly = plg.Polygon(box)
|
|
area = poly.area()
|
|
if area < 10:
|
|
continue
|
|
box = np.array(box)
|
|
boxes.append(box)
|
|
return np.array(boxes)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
image = cv2.imread('images/standard.jpg', cv2.IMREAD_COLOR)
|
|
boxes = watershed(image, True)
|
|
print(boxes) |