113 lines
4.2 KiB
Python
113 lines
4.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; will use whole-image 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):
|
|
if not self.east_available:
|
|
h, w = img_rgb.shape[:2]
|
|
return [(_int0(0), _int0(0), _int0(w), _int0(h))], False
|
|
|
|
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)))
|
|
if not boxes:
|
|
boxes.append((_int0(0), _int0(0), _int0(w), _int0(h)))
|
|
return boxes, True
|
|
|
|
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 _int0(val):
|
|
try:
|
|
return int(val)
|
|
except Exception:
|
|
return 0
|