46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
import cv2
|
|
import numpy as np
|
|
from skimage.segmentation import watershed
|
|
|
|
|
|
def segment_region_score(watershed_param, region_score, word_image, pseudo_vis_opt):
|
|
region_score = np.float32(region_score) / 255
|
|
fore = np.uint8(region_score > 0.75)
|
|
back = np.uint8(region_score < 0.05)
|
|
unknown = 1 - (fore + back)
|
|
ret, markers = cv2.connectedComponents(fore)
|
|
markers += 1
|
|
markers[unknown == 1] = 0
|
|
|
|
labels = watershed(-region_score, markers)
|
|
boxes = []
|
|
for label in range(2, ret + 1):
|
|
y, x = np.where(labels == label)
|
|
x_max = x.max()
|
|
y_max = y.max()
|
|
x_min = x.min()
|
|
y_min = y.min()
|
|
box = [[x_min, y_min], [x_max, y_min], [x_max, y_max], [x_min, y_max]]
|
|
box = np.array(box)
|
|
box *= 2
|
|
boxes.append(box)
|
|
return np.array(boxes, dtype=np.float32)
|
|
|
|
|
|
def exec_watershed_by_version(
|
|
watershed_param, region_score, word_image, pseudo_vis_opt
|
|
):
|
|
|
|
func_name_map_dict = {
|
|
"skimage": segment_region_score,
|
|
}
|
|
|
|
try:
|
|
return func_name_map_dict[watershed_param.version](
|
|
watershed_param, region_score, word_image, pseudo_vis_opt
|
|
)
|
|
except:
|
|
print(
|
|
f"Watershed version {watershed_param.version} does not exist in func_name_map_dict."
|
|
)
|