前段时间在跑文本检测的psenet模型，psenet的后处理过程使用了使用了一个称为PSE(progressive scale expansion，逐步的尺度扩张)的处理步骤来得到完整的word bbox，作者提供了C++和Python的PSE实现，其中使用Python版本的PSE非常缓慢。
最近在复现另一个文本检测模型CRAFT的过程中，接触到了用于分割的watershed/分水岭算法，opencv提供了watershed的函数接口cv2.watershed()。
经过简单了解之后，发现分水岭算法的原理其实和PSE差不多，都是通过一些最初指定的kernel，然后不断向外扩张来达到图像分割的效果。于是试着用cv2.watershed()来代替Python版本的PSE处理。
相关代码如下：
...
img = img.resize((resize_w, resize_h), Image.BILINEAR)
input_img = transform(img).unsqueeze(0).to(device)
with torch.no_grad():
outputs = model(input_img)
outputs = torch.sigmoid(outputs)
score = outputs[:, 0, :, :]
outputs = outputs > args.threshold # torch.uint8
text = outputs[:, 0, :, :]
kernels = outputs[:, 0:args.kernel_num, :, :] * text
score = score.squeeze(0).cpu().numpy()
text = text.squeeze(0).cpu().numpy()
kernels = kernels.squeeze(0).cpu().numpy()
tmp_marker = kernels[-1, :, :]
for i in range(args.kernel_num-2, -1, -1):
sure_fg = tmp_marker
sure_bg = kernels[i, :, :]
watershed_source = cv2.cvtColor(sure_bg, cv2.COLOR_GRAY2BGR)
unknown = cv2.subtract(sure_bg,sure_fg)
ret, marker = cv2.connectedComponents(sure_fg)
label_num = np.max(marker)
marker += 1
marker[unknown==1] = 0
marker = cv2.watershed(watershed_source, marker)
marker[marker==-1] = 1
marker -= 1
tmp_marker = np.asarray(marker, np.uint8)
label = tmp_marker
scale = (w / marker.shape[1], h / marker.shape[0])
bboxes = []
for i in range(1, label_num+1):
# get [x,y] pair, points.shape=[n, 2]
points = np.array(np.where(label == i)).transpose((1, 0))[:, ::-1]
# similar to pixellink's min_area when post-processing
if points.shape[0] < args.min_area / (args.scale * args.scale):
continue
#this filter op is very important, f-score=68.0(without) vs 69.1(with)
score_i = np.mean(score[label == i])
if score_i < args.min_score:
continue
rect = cv2.minAreaRect(points)
bbox = cv2.boxPoints(rect) * scale
bbox = bbox.astype('int32')
bboxes.append(bbox.reshape(-1))
...
上面的代码模拟了PSE的过程，在ic15测试集上跑，速度比c++版本的PSE还快一些(使用的resnet152，速度对比：1.28fps vs 1.05fps)。但是准确率下降了约3个点(f-score：82.3 vs 85.4)。
不过我直接使用最大尺度的kernel作为watershed的源，不使用PSE的过程，直接从最小scale的kernel扩张到最大尺度的kernel，这样得到的结果反而更好一些，f-score达到了84.2。虽然比作者提供的PSE算法低一些，但是速度更快，能达到1.42fps，而且代码也挺简单的。
修改后的部分代码：
...
sure_fg = kernels[-1, :, :]
sure_bg = text
watershed_source = cv2.cvtColor(sure_bg, cv2.COLOR_GRAY2BGR)
unknown = cv2.subtract(sure_bg,sure_fg)
ret, marker = cv2.connectedComponents(sure_fg)
label_num = np.max(marker)
marker += 1
marker[unknown==1] = 0
marker = cv2.watershed(watershed_source, marker)
marker -= 1
label = marker
...
参考资料：
- Image Segmentation with Watershed Algorithm
- How to define the markers for Watershed in OpenCV?
- PSENet