前段時(shí)間在跑文本檢測(cè)的psenet模型秃臣，psenet的后處理過(guò)程使用了使用了一個(gè)稱為PSE(progressive scale expansion塞俱，逐步的尺度擴(kuò)張)的處理步驟來(lái)得到完整的word bbox蜘渣，作者提供了C++和Python的PSE實(shí)現(xiàn)授舟，其中使用Python版本的PSE非常緩慢塞弊。
最近在復(fù)現(xiàn)另一個(gè)文本檢測(cè)模型CRAFT的過(guò)程中肩碟，接觸到了用于分割的watershed/分水嶺算法强窖，opencv提供了watershed的函數(shù)接口cv2.watershed()。
經(jīng)過(guò)簡(jiǎn)單了解之后削祈，發(fā)現(xiàn)分水嶺算法的原理其實(shí)和PSE差不多翅溺，都是通過(guò)一些最初指定的kernel脑漫，然后不斷向外擴(kuò)張來(lái)達(dá)到圖像分割的效果。于是試著用cv2.watershed()來(lái)代替Python版本的PSE處理咙崎。
相關(guān)代碼如下：
...
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的過(guò)程优幸，在ic15測(cè)試集上跑，速度比c++版本的PSE還快一些(使用的resnet152褪猛，速度對(duì)比：1.28fps vs 1.05fps)网杆。但是準(zhǔn)確率下降了約3個(gè)點(diǎn)(f-score：82.3 vs 85.4)。
不過(guò)我直接使用最大尺度的kernel作為watershed的源伊滋，不使用PSE的過(guò)程碳却，直接從最小scale的kernel擴(kuò)張到最大尺度的kernel，這樣得到的結(jié)果反而更好一些笑旺，f-score達(dá)到了84.2昼浦。雖然比作者提供的PSE算法低一些，但是速度更快筒主，能達(dá)到1.42fps关噪，而且代碼也挺簡(jiǎn)單的。
修改后的部分代碼：
...
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
...
參考資料：
