256 lines
9.2 KiB
Python
256 lines
9.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Text region detection using OpenCV EAST with graceful fallback."""
|
|
import os
|
|
import sys
|
|
import numpy as np
|
|
|
|
try:
|
|
import cv2
|
|
except ImportError: # pragma: no cover
|
|
cv2 = None
|
|
|
|
|
|
class TextDetector(object):
|
|
def __init__(self, east_model_path, score_threshold=0.4, use_cuda=False):
|
|
self.east_model_path = east_model_path
|
|
self.score_threshold = score_threshold
|
|
self.use_cuda = use_cuda
|
|
self.net = None
|
|
self.east_available = False
|
|
self._load()
|
|
|
|
def _load(self):
|
|
if cv2 is None:
|
|
print('[WARN] OpenCV not installed; EAST disabled', file=sys.stderr)
|
|
return
|
|
if not os.path.exists(self.east_model_path):
|
|
print('[INFO] EAST model not found; using OpenCV contour fallback', file=sys.stderr)
|
|
return
|
|
try:
|
|
self.net = cv2.dnn.readNet(self.east_model_path)
|
|
if self.use_cuda and hasattr(cv2, "cuda") and cv2.cuda.getCudaEnabledDeviceCount() > 0:
|
|
try:
|
|
self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
|
|
self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
|
|
print('[INFO] EAST using CUDA backend/target', file=sys.stderr)
|
|
except Exception as exc:
|
|
print('[WARN] Failed to enable CUDA for EAST: {}'.format(exc), file=sys.stderr)
|
|
self.east_available = True
|
|
print('[INFO] EAST model loaded: {}'.format(self.east_model_path))
|
|
except Exception as exc:
|
|
print('[WARN] Failed to load EAST model: {}'.format(exc), file=sys.stderr)
|
|
self.east_available = False
|
|
|
|
def detect(self, img_rgb):
|
|
h, w = img_rgb.shape[:2]
|
|
if self.east_available:
|
|
boxes = self._detect_with_east(img_rgb)
|
|
if boxes:
|
|
return boxes, True
|
|
|
|
boxes = self._detect_with_cv_fallback(img_rgb)
|
|
if boxes:
|
|
return boxes, False
|
|
|
|
return [(_int0(0), _int0(0), _int0(w), _int0(h))], False
|
|
|
|
def _detect_with_east(self, img_rgb):
|
|
h, w = img_rgb.shape[:2]
|
|
# EAST expects width/height divisible by 32
|
|
new_w = 320
|
|
new_h = 320
|
|
blob = cv2.dnn.blobFromImage(img_rgb, 1.0, (new_w, new_h),
|
|
(123.68, 116.78, 103.94), swapRB=True, crop=False)
|
|
self.net.setInput(blob)
|
|
scores, geometry = self.net.forward([
|
|
"feature_fusion/Conv_7/Sigmoid",
|
|
"feature_fusion/concat_3"
|
|
])
|
|
rectangles, confidences = self._decode(scores, geometry, self.score_threshold)
|
|
indices = cv2.dnn.NMSBoxes(rectangles, confidences, self.score_threshold, 0.4)
|
|
|
|
boxes = []
|
|
rW = float(w) / float(new_w)
|
|
rH = float(h) / float(new_h)
|
|
if len(indices) > 0:
|
|
for i in indices.flatten():
|
|
x, y, bw, bh = rectangles[i]
|
|
x1 = int(x * rW)
|
|
y1 = int(y * rH)
|
|
x2 = int((x + bw) * rW)
|
|
y2 = int((y + bh) * rH)
|
|
boxes.append((_int0(x1), _int0(y1), _int0(x2), _int0(y2)))
|
|
return self._normalize_boxes(boxes, w, h, max_boxes=16)
|
|
|
|
def _decode(self, scores, geometry, score_thresh):
|
|
num_rows, num_cols = scores.shape[2:4]
|
|
rectangles = []
|
|
confidences = []
|
|
for y in range(num_rows):
|
|
scores_data = scores[0, 0, y]
|
|
x0 = geometry[0, 0, y]
|
|
x1 = geometry[0, 1, y]
|
|
x2 = geometry[0, 2, y]
|
|
x3 = geometry[0, 3, y]
|
|
angles = geometry[0, 4, y]
|
|
for x in range(num_cols):
|
|
score = scores_data[x]
|
|
if score < score_thresh:
|
|
continue
|
|
offset_x = x * 4.0
|
|
offset_y = y * 4.0
|
|
angle = angles[x]
|
|
cos = np.cos(angle)
|
|
sin = np.sin(angle)
|
|
h = x0[x] + x2[x]
|
|
w = x1[x] + x3[x]
|
|
end_x = int(offset_x + (cos * x1[x]) + (sin * x2[x]))
|
|
end_y = int(offset_y - (sin * x1[x]) + (cos * x2[x]))
|
|
start_x = int(end_x - w)
|
|
start_y = int(end_y - h)
|
|
rectangles.append((start_x, start_y, int(w), int(h)))
|
|
confidences.append(float(score))
|
|
return rectangles, confidences
|
|
|
|
def _detect_with_cv_fallback(self, img_rgb):
|
|
if cv2 is None:
|
|
return []
|
|
try:
|
|
gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
|
|
except Exception:
|
|
return []
|
|
|
|
h, w = gray.shape[:2]
|
|
img_area = float(max(1, h * w))
|
|
min_area = max(160.0, img_area * 0.00035)
|
|
|
|
boxes = []
|
|
|
|
# Pass #1: gradient + morphology for horizontal text lines.
|
|
try:
|
|
grad_x = cv2.Sobel(gray, ddepth=cv2.CV_32F, dx=1, dy=0, ksize=3)
|
|
grad_x = cv2.convertScaleAbs(grad_x)
|
|
grad_x = cv2.GaussianBlur(grad_x, (3, 3), 0)
|
|
_, bin1 = cv2.threshold(grad_x, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
|
k_close = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 5))
|
|
k_dilate = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 3))
|
|
bin1 = cv2.morphologyEx(bin1, cv2.MORPH_CLOSE, k_close, iterations=1)
|
|
bin1 = cv2.dilate(bin1, k_dilate, iterations=1)
|
|
boxes.extend(self._boxes_from_mask(bin1, w, h, min_area=min_area))
|
|
except Exception:
|
|
pass
|
|
|
|
# Pass #2: adaptive threshold for low-contrast / noisy scenes.
|
|
try:
|
|
bin2 = cv2.adaptiveThreshold(
|
|
gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 31, 15
|
|
)
|
|
k_close = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 3))
|
|
k_dilate = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
|
|
bin2 = cv2.morphologyEx(bin2, cv2.MORPH_CLOSE, k_close, iterations=1)
|
|
bin2 = cv2.dilate(bin2, k_dilate, iterations=1)
|
|
boxes.extend(self._boxes_from_mask(bin2, w, h, min_area=min_area))
|
|
except Exception:
|
|
pass
|
|
|
|
return self._normalize_boxes(boxes, w, h, max_boxes=18)
|
|
|
|
def _boxes_from_mask(self, mask, image_w, image_h, min_area):
|
|
out = []
|
|
contours_result = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
contours = contours_result[0] if len(contours_result) == 2 else contours_result[1]
|
|
min_h = max(10, int(image_h * 0.015))
|
|
min_w = max(16, int(image_w * 0.02))
|
|
for cnt in contours:
|
|
x, y, bw, bh = cv2.boundingRect(cnt)
|
|
area = float(bw * bh)
|
|
if area < min_area:
|
|
continue
|
|
if bw < min_w or bh < min_h:
|
|
continue
|
|
aspect = float(bw) / float(max(1, bh))
|
|
if aspect < 0.7 or aspect > 45.0:
|
|
continue
|
|
contour_area = float(cv2.contourArea(cnt))
|
|
fill_ratio = contour_area / float(max(1, bw * bh))
|
|
if fill_ratio < 0.03:
|
|
continue
|
|
pad_x = max(2, int(bw * 0.08))
|
|
pad_y = max(2, int(bh * 0.30))
|
|
x1 = max(0, x - pad_x)
|
|
y1 = max(0, y - pad_y)
|
|
x2 = min(image_w, x + bw + pad_x)
|
|
y2 = min(image_h, y + bh + pad_y)
|
|
out.append((_int0(x1), _int0(y1), _int0(x2), _int0(y2)))
|
|
return out
|
|
|
|
def _normalize_boxes(self, boxes, image_w, image_h, max_boxes=16):
|
|
norm = []
|
|
for box in boxes or []:
|
|
x1, y1, x2, y2 = _clip_box(box, image_w, image_h)
|
|
if x2 <= x1 or y2 <= y1:
|
|
continue
|
|
norm.append((x1, y1, x2, y2))
|
|
if not norm:
|
|
return []
|
|
kept = _nms_by_iou(norm, iou_thresh=0.45)
|
|
kept = sorted(kept, key=lambda b: ((b[3] - b[1]) * (b[2] - b[0])), reverse=True)
|
|
if max_boxes and len(kept) > max_boxes:
|
|
kept = kept[:max_boxes]
|
|
kept = sorted(kept, key=lambda b: (b[1], b[0]))
|
|
return kept
|
|
|
|
|
|
def _int0(val):
|
|
try:
|
|
return int(val)
|
|
except Exception:
|
|
return 0
|
|
|
|
|
|
def _clip_box(box, image_w, image_h):
|
|
x1, y1, x2, y2 = box
|
|
x1 = max(0, min(_int0(x1), _int0(image_w)))
|
|
y1 = max(0, min(_int0(y1), _int0(image_h)))
|
|
x2 = max(0, min(_int0(x2), _int0(image_w)))
|
|
y2 = max(0, min(_int0(y2), _int0(image_h)))
|
|
return x1, y1, x2, y2
|
|
|
|
|
|
def _nms_by_iou(boxes, iou_thresh=0.45):
|
|
if not boxes:
|
|
return []
|
|
ranked = sorted(boxes, key=_box_area, reverse=True)
|
|
keep = []
|
|
for candidate in ranked:
|
|
should_keep = True
|
|
for chosen in keep:
|
|
if _iou(candidate, chosen) >= iou_thresh:
|
|
should_keep = False
|
|
break
|
|
if should_keep:
|
|
keep.append(candidate)
|
|
return keep
|
|
|
|
|
|
def _box_area(box):
|
|
x1, y1, x2, y2 = box
|
|
return max(0, x2 - x1) * max(0, y2 - y1)
|
|
|
|
|
|
def _iou(a, b):
|
|
ax1, ay1, ax2, ay2 = a
|
|
bx1, by1, bx2, by2 = b
|
|
xx1 = max(ax1, bx1)
|
|
yy1 = max(ay1, by1)
|
|
xx2 = min(ax2, bx2)
|
|
yy2 = min(ay2, by2)
|
|
iw = max(0, xx2 - xx1)
|
|
ih = max(0, yy2 - yy1)
|
|
inter = iw * ih
|
|
if inter <= 0:
|
|
return 0.0
|
|
union = float(_box_area(a) + _box_area(b) - inter)
|
|
if union <= 0:
|
|
return 0.0
|
|
return inter / union
|