ScreenSubtitleReader 0.1.3__py3-none-any.whl

ScreenSubtitleReader/ocr_module/dbnet/decode.py

@@ -0,0 +1,125 @@
+import cv2
+import numpy as np
+import pyclipper
+from shapely.geometry import Polygon
+
+
+class SegDetectorRepresenter:
+    def __init__(self, thresh=0.3, box_thresh=0.5, max_candidates=1000, unclip_ratio=2.0):
+        self.min_size = 3
+        self.thresh = thresh
+        self.box_thresh = box_thresh
+        self.max_candidates = max_candidates
+        self.unclip_ratio = unclip_ratio
+    
+    def __call__(self, pred, height, width):
+        """
+        batch: (image, polygons, ignore_tags
+        batch: a dict produced by dataloaders.
+            image: tensor of shape (N, C, H, W).
+            polygons: tensor of shape (N, K, 4, 2), the polygons of objective regions.
+            ignore_tags: tensor of shape (N, K), indicates whether a region is ignorable or not.
+            shape: the original shape of images.
+            filename: the original filenames of images.
+        pred:
+            binary: text region segmentation map, with shape (N, H, W)
+            thresh: [if exists] thresh hold prediction with shape (N, H, W)
+            thresh_binary: [if exists] binarized with threshhold, (N, H, W)
+        """
+
+        pred = pred[0, :, :]
+        segmentation = self.binarize(pred)
+        
+        boxes, scores = self.boxes_from_bitmap(pred, segmentation, width, height)
+        
+        return boxes, scores
+    
+    def binarize(self, pred):
+        return pred > self.thresh
+    
+    def boxes_from_bitmap(self, pred, bitmap, dest_width, dest_height):
+        """
+        _bitmap: single map with shape (H, W),
+            whose values are binarized as {0, 1}
+        """
+        
+        assert len(bitmap.shape) == 2
+        # bitmap = _bitmap.cpu().numpy()  # The first channel
+        # pred = pred.cpu().detach().numpy()
+        height, width = bitmap.shape
+        # print(bitmap)
+        # cv2.imwrite("./test/output/test.jpg",(bitmap * 255).astype(np.uint8))
+        contours, _ = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+        num_contours = min(len(contours), self.max_candidates)
+        boxes = np.zeros((num_contours, 4, 2), dtype=np.int16)
+        scores = np.zeros((num_contours,), dtype=np.float32)
+        rects = []
+        for index in range(num_contours):
+            contour = contours[index].squeeze(1)
+            points, sside = self.get_mini_boxes(contour)
+            if sside < self.min_size:
+                continue
+            points = np.array(points)
+            score = self.box_score_fast(pred, contour)
+            if self.box_thresh > score:
+                continue
+            # print(points)
+            box = self.unclip(points, unclip_ratio=self.unclip_ratio).reshape(-1, 1, 2)
+            # print(box)
+            box, sside = self.get_mini_boxes(box)
+            if sside < self.min_size + 2:
+                continue
+            box = np.array(box)
+            if not isinstance(dest_width, int):
+                dest_width = dest_width.item()
+                dest_height = dest_height.item()
+            
+            box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
+            box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
+            boxes[index, :, :] = box.astype(np.int16)
+            scores[index] = score
+        return boxes, scores
+    
+    def unclip(self, box, unclip_ratio=1.5):
+        poly = Polygon(box)
+        
+        distance = poly.area * unclip_ratio / (poly.length  )
+        offset = pyclipper.PyclipperOffset()
+        offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+        expanded = np.array(offset.Execute(distance))
+        return expanded
+    
+    def get_mini_boxes(self, contour):
+        bounding_box = cv2.minAreaRect(contour)
+        points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
+        
+        index_1, index_2, index_3, index_4 = 0, 1, 2, 3
+        if points[1][1] > points[0][1]:
+            index_1 = 0
+            index_4 = 1
+        else:
+            index_1 = 1
+            index_4 = 0
+        if points[3][1] > points[2][1]:
+            index_2 = 2
+            index_3 = 3
+        else:
+            index_2 = 3
+            index_3 = 2
+        
+        box = [points[index_1], points[index_2], points[index_3], points[index_4]]
+        return box, min(bounding_box[1])
+    
+    def box_score_fast(self, bitmap, _box):
+        h, w = bitmap.shape[:2]
+        box = _box.copy()
+        xmin = np.clip(np.floor(box[:, 0].min()).astype(int), 0, w - 1)
+        xmax = np.clip(np.ceil(box[:, 0].max()).astype(int), 0, w - 1)
+        ymin = np.clip(np.floor(box[:, 1].min()).astype(int), 0, h - 1)
+        ymax = np.clip(np.ceil(box[:, 1].max()).astype(int), 0, h - 1)
+        
+        mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
+        box[:, 0] = box[:, 0] - xmin
+        box[:, 1] = box[:, 1] - ymin
+        cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)
+        return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]

ScreenSubtitleReader/ocr_module/model.py

@@ -0,0 +1,128 @@
+from .config import *
+from .crnn import CRNNHandle
+from .angnet import  AngleNetHandle
+from .utils import draw_bbox, crop_rect, sorted_boxes, get_rotate_crop_image
+from PIL import Image
+import numpy as np
+import cv2
+import copy
+from .dbnet.dbnet_infer import DBNET
+import time
+import traceback
+
+class  OcrHandle(object):
+    def __init__(self):
+        self.text_handle = DBNET(model_path)
+        self.crnn_handle = CRNNHandle(crnn_model_path)
+        if angle_detect:
+            self.angle_handle = AngleNetHandle(angle_net_path)
+
+
+    def crnnRecWithBox(self,im, boxes_list,score_list):
+        """
+        crnn模型，ocr识别
+        @@model,
+        @@converter,
+        @@im:Array
+        @@text_recs:text box
+        @@ifIm:是否输出box对应的img
+
+        """
+        results = []
+        boxes_list = sorted_boxes(np.array(boxes_list))
+
+        line_imgs = []
+        for index, (box, score) in enumerate(zip(boxes_list[:angle_detect_num], score_list[:angle_detect_num])):
+            tmp_box = copy.deepcopy(box)
+            partImg_array = get_rotate_crop_image(im, tmp_box.astype(np.float32))
+            partImg = Image.fromarray(partImg_array).convert("RGB")
+            line_imgs.append(partImg)
+
+        angle_res = False
+        if angle_detect:
+            angle_res = self.angle_handle.predict_rbgs(line_imgs)
+
+        count = 1
+        for index, (box ,score) in enumerate(zip(boxes_list,score_list)):
+
+            tmp_box = copy.deepcopy(box)
+            partImg_array = get_rotate_crop_image(im, tmp_box.astype(np.float32))
+
+
+            partImg = Image.fromarray(partImg_array).convert("RGB")
+
+            if angle_detect and angle_res:
+                partImg = partImg.rotate(180)
+
+
+            if not is_rgb:
+                partImg = partImg.convert('L')
+
+            try:
+                if is_rgb:
+                    simPred = self.crnn_handle.predict_rbg(partImg)  ##识别的文本
+                else:
+                    simPred = self.crnn_handle.predict(partImg)  ##识别的文本
+            except Exception as e:
+                print(traceback.format_exc())
+                continue
+
+            if simPred.strip() != '':
+                results.append([tmp_box,simPred,score])
+                count += 1
+
+        return results
+
+
+    def text_predict(self,img,short_size):
+        boxes_list, score_list = self.text_handle.process(np.asarray(img).astype(np.uint8),short_size=short_size)
+        result = self.crnnRecWithBox(np.array(img), boxes_list,score_list)
+
+        return result
+
+
+if __name__ == "__main__":
+    from PIL import ImageGrab
+
+
+    print("正在初始化 OCR 模型，请稍候...")
+    try:
+        ocr = OcrHandle()
+        print("✅ 模型初始化成功！")
+    except Exception as e:
+        print(f"❌ 模型初始化失败。\n错误信息: {e}")
+        exit(1)
+
+    print("📸 正在截取屏幕...")
+    # 使用 ImageGrab 截取全屏
+    im = ImageGrab.grab()
+    
+    # 【关键修正】颜色通道转换：
+    # ImageGrab 截取的是 RGB 格式的 PIL 图像，而原项目默认使用 cv2.imread (BGR 格式)。
+    # 为了保证模型识别效果不出偏差，这里将其转为 OpenCV 标准的 BGR numpy 数组
+    img = cv2.cvtColor(np.array(im), cv2.COLOR_RGB2BGR)
+    
+    print("🔍 开始执行 OCR 识别...")
+    start_time = time.time()
+    
+    # 调用识别接口 (如果是全屏高分辨率截图，short_size 可以保持 960，或者视情况调高)
+    results = ocr.text_predict(img, short_size=960)
+    
+    cost_time = time.time() - start_time
+    print(f"✅ 识别完成！耗时: {cost_time:.3f} 秒")
+    print("=" * 50)
+    
+    if not results:
+        print("未能识别到任何文字。")
+    else:
+        print("部分识别结果如下 (最多显示前20条):")
+        # 截屏识别出来的文本通常很多，测试时为了清爽只打印前 20 条
+        for i, res in enumerate(results[:20]):
+            text = res[1] 
+            score = res[2]
+            print(f"文本: {text.strip()}  |  置信度: {score:.4f}")
+            
+        if len(results) > 20:
+            print(f"... (还有 {len(results) - 20} 条未显示)")
+            
+    print("=" * 50)

ScreenSubtitleReader/ocr_module/models/__init__.py

@@ -0,0 +1 @@
+"""Packaged ONNX model files for chineseocr_lite."""

ScreenSubtitleReader/ocr_module/utils.py

@@ -0,0 +1,173 @@
+import numpy as np
+import cv2
+from PIL import Image
+
+
+def rotate_cut_img(im, degree, x_center, y_center, w, h, leftAdjust=False, rightAdjust=False, alph=0.2):
+    # degree_ = degree * 180.0 / np.pi
+    # print(degree_)
+    right = 0
+    left = 0
+    if rightAdjust:
+        right = 1
+    if leftAdjust:
+        left = 1
+    
+    box = (max(1, x_center - w / 2 - left * alph * (w / 2))
+           , y_center - h / 2,  # ymin
+           min(x_center + w / 2 + right * alph * (w / 2), im.size[0] - 1)
+           , y_center + h / 2)  # ymax
+    
+    newW = box[2] - box[0]
+    newH = box[3] - box[1]
+    tmpImg = im.rotate(degree, center=(x_center, y_center)).crop(box)
+    
+    return tmpImg, newW, newH
+
+
+def crop_rect(img, rect, alph=0.15):
+    img = np.asarray(img)
+    # get the parameter of the small rectangle
+    # print("rect!")
+    # print(rect)
+    center, size, angle = rect[0], rect[1], rect[2]
+    min_size = min(size)
+    
+    if angle > -45:
+        center, size = tuple(map(int, center)), tuple(map(int, size))
+        # angle-=270
+        size = (int(size[0] + min_size * alph), int(size[1] + min_size * alph))
+        height, width = img.shape[0], img.shape[1]
+        M = cv2.getRotationMatrix2D(center, angle, 1)
+        # size = tuple([int(rect[1][1]), int(rect[1][0])])
+        img_rot = cv2.warpAffine(img, M, (width, height))
+        # cv2.imwrite("debug_im/img_rot.jpg", img_rot)
+        img_crop = cv2.getRectSubPix(img_rot, size, center)
+    else:
+        center = tuple(map(int, center))
+        size = tuple([int(rect[1][1]), int(rect[1][0])])
+        size = (int(size[0] + min_size * alph), int(size[1] + min_size * alph))
+        angle -= 270
+        height, width = img.shape[0], img.shape[1]
+        M = cv2.getRotationMatrix2D(center, angle, 1)
+        img_rot = cv2.warpAffine(img, M, (width, height))
+        # cv2.imwrite("debug_im/img_rot.jpg", img_rot)
+        img_crop = cv2.getRectSubPix(img_rot, size, center)
+    img_crop = Image.fromarray(img_crop)
+    return img_crop
+
+
+def draw_bbox(img_path, result, color=(255, 0, 0), thickness=2):
+    if isinstance(img_path, str):
+        img_path = cv2.imread(img_path)
+        # img_path = cv2.cvtColor(img_path, cv2.COLOR_BGR2RGB)
+    img_path = img_path.copy()
+    for point in result:
+        point = point.astype(int)
+        cv2.line(img_path, tuple(point[0]), tuple(point[1]), color, thickness)
+        cv2.line(img_path, tuple(point[1]), tuple(point[2]), color, thickness)
+        cv2.line(img_path, tuple(point[2]), tuple(point[3]), color, thickness)
+        cv2.line(img_path, tuple(point[3]), tuple(point[0]), color, thickness)
+    return img_path
+
+
+def sort_box(boxs):
+    res = []
+    for box in boxs:
+        # box = [x if x>0 else 0 for x in box ]
+        x1, y1, x2, y2, x3, y3, x4, y4 = box[:8]
+        newBox = [[x1, y1], [x2, y2], [x3, y3], [x4, y4]]
+        # sort x
+        newBox = sorted(newBox, key=lambda x: x[0])
+        x1, y1 = sorted(newBox[:2], key=lambda x: x[1])[0]
+        index = newBox.index([x1, y1])
+        newBox.pop(index)
+        newBox = sorted(newBox, key=lambda x: -x[1])
+        x4, y4 = sorted(newBox[:2], key=lambda x: x[0])[0]
+        index = newBox.index([x4, y4])
+        newBox.pop(index)
+        newBox = sorted(newBox, key=lambda x: -x[0])
+        x2, y2 = sorted(newBox[:2], key=lambda x: x[1])[0]
+        index = newBox.index([x2, y2])
+        newBox.pop(index)
+        
+        newBox = sorted(newBox, key=lambda x: -x[1])
+        x3, y3 = sorted(newBox[:2], key=lambda x: x[0])[0]
+        
+        res.append([x1, y1, x2, y2, x3, y3, x4, y4])
+    return res
+
+
+def solve(box):
+    """
+    绕 cx,cy点 w,h 旋转 angle 的坐标
+    x = cx-w/2
+    y = cy-h/2
+    x1-cx = -w/2*cos(angle) +h/2*sin(angle)
+    y1 -cy= -w/2*sin(angle) -h/2*cos(angle)
+
+    h(x1-cx) = -wh/2*cos(angle) +hh/2*sin(angle)
+    w(y1 -cy)= -ww/2*sin(angle) -hw/2*cos(angle)
+    (hh+ww)/2sin(angle) = h(x1-cx)-w(y1 -cy)
+
+    """
+    x1, y1, x2, y2, x3, y3, x4, y4 = box[:8]
+    cx = (x1 + x3 + x2 + x4) / 4.0
+    cy = (y1 + y3 + y4 + y2) / 4.0
+    w = (np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2) + np.sqrt((x3 - x4) ** 2 + (y3 - y4) ** 2)) / 2
+    h = (np.sqrt((x2 - x3) ** 2 + (y2 - y3) ** 2) + np.sqrt((x1 - x4) ** 2 + (y1 - y4) ** 2)) / 2
+    
+    sinA = (h * (x1 - cx) - w * (y1 - cy)) * 1.0 / (h * h + w * w) * 2
+    angle = np.arcsin(sinA)
+    return angle, w, h, cx, cy
+
+
+def sorted_boxes(dt_boxes):
+    """
+    Sort text boxes in order from top to bottom, left to right
+    args:
+        dt_boxes(array):detected text boxes with shape [4, 2]
+    return:
+        sorted boxes(array) with shape [4, 2]
+    """
+    num_boxes = dt_boxes.shape[0]
+    sorted_boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
+    _boxes = list(sorted_boxes)
+
+    for i in range(num_boxes - 1):
+        if abs(_boxes[i+1][0][1] - _boxes[i][0][1]) < 10 and \
+            (_boxes[i + 1][0][0] < _boxes[i][0][0]):
+            tmp = _boxes[i]
+            _boxes[i] = _boxes[i + 1]
+            _boxes[i + 1] = tmp
+    return _boxes
+
+
+def get_rotate_crop_image(img, points):
+    img_height, img_width = img.shape[0:2]
+    left = int(np.min(points[:, 0]))
+    right = int(np.max(points[:, 0]))
+    top = int(np.min(points[:, 1]))
+    bottom = int(np.max(points[:, 1]))
+    img_crop = img[top:bottom, left:right, :].copy()
+    points[:, 0] = points[:, 0] - left
+    points[:, 1] = points[:, 1] - top
+    img_crop_width = int(np.linalg.norm(points[0] - points[1]))
+    img_crop_height = int(np.linalg.norm(points[0] - points[3]))
+    pts_std = np.float32([[0, 0], [img_crop_width, 0],\
+        [img_crop_width, img_crop_height], [0, img_crop_height]])
+
+    M = cv2.getPerspectiveTransform(points, pts_std)
+    dst_img = cv2.warpPerspective(
+        img_crop,
+        M, (img_crop_width, img_crop_height),
+        borderMode=cv2.BORDER_REPLICATE)
+    dst_img_height, dst_img_width = dst_img.shape[0:2]
+    if dst_img_height * 1.0 / dst_img_width >= 1.5:
+        dst_img = np.rot90(dst_img)
+    return dst_img
+
+
+def app_url(version, name):
+    url = '/{}/{}'.format(version, name)
+    return url

ScreenSubtitleReader/setting/setting.py

@@ -0,0 +1,98 @@
+import json
+import threading
+from pathlib import Path
+# 如果尚未安装，可以运行: pip install platformdirs
+from platformdirs import user_data_dir
+
+class AppConfig:
+    _instance = None
+    _lock = threading.Lock()  # 确保线程安全的单例模式
+
+    def __new__(cls, *args, **kwargs):
+        """单例模式：确保全局只有一个配置实例"""
+        if not cls._instance:
+            with cls._lock:
+                if not cls._instance:
+                    cls._instance = super().__new__(cls)
+                    # 标记是否已经初始化，防止 __init__ 重复调用
+                    cls._instance._initialized = False
+        return cls._instance
+
+    def __init__(self, app_name="ScreenSubtitleReader"):
+        if self._initialized:
+            return
+        
+        # 1. 确定跨平台的 AppData 存储路径
+        # 在 Windows 上通常是: C:\Users\用户名\AppData\Local\MyTkinterApp
+        self.config_dir = Path(user_data_dir(app_name))
+        self.config_file = self.config_dir / "config.json"
+        
+        # 2. 默认属性
+        self._volume = 100
+        self._speed = 270
+        
+        # 3. 启动初始化：自动读取或创建文件
+        self._load_or_create_config()
+        self._initialized = True
+
+    def _load_or_create_config(self):
+        """读取配置文件，如果不存在则创建默认配置"""
+        try:
+            if self.config_file.exists():
+                # 文件存在，读取数据
+                with open(self.config_file, "r", encoding="utf-8") as f:
+                    data = json.load(f)
+                    # 使用 get 提供降级容错，防止 JSON 被人为改坏导致崩溃
+                    self._volume = int(data.get("volume", 100))
+                    self._speed = int(data.get("speed", 270))
+                print(f"成功加载配置: {self.config_file}")
+            else:
+                # 文件不存在，保存当前的默认值
+                self.save()
+                print(f"创建默认配置文件: {self.config_file}")
+        except Exception as e:
+            print(f"配置加载或初始化失败，使用默认值。错误: {e}")
+
+    def save(self):
+        """将当前属性序列化并自动存盘"""
+        try:
+            # 确保应用数据文件夹存在（不存在则自动递归创建）
+            self.config_dir.mkdir(parents=True, exist_ok=True)
+            
+            data = {
+                "volume": self._volume,
+                "speed": self._speed
+            }
+            # 写入文件
+            with open(self.config_file, "w", encoding="utf-8") as f:
+                json.dump(data, f, indent=4, ensure_ascii=False)
+            print("配置已自动保存到磁盘。")
+        except Exception as e:
+            print(f"配置保存失败: {e}")
+
+    # --- 使用 Property 属性装饰器，既能像属性一样访问，又能触发自动存盘 ---
+    @property
+    def volume(self) -> int:
+        return self._volume
+
+    @volume.setter
+    def volume(self, value: int):
+        if not isinstance(value, int):
+            raise TypeError("音量必须是整型 (int)")
+        self._volume = value
+        self.save()  # 更改时自动存盘
+
+    @property
+    def speed(self) -> int:
+        return self._speed
+
+    @speed.setter
+    def speed(self, value: int):
+        if not isinstance(value, int):
+            raise TypeError("语速必须是整型 (int)")
+        self._speed = value
+        self.save()  # 更改时自动存盘
+        
+if __name__ == "__main__":
+    apc = AppConfig()
+    print(apc.speed)

ScreenSubtitleReader/speaker/speaker.py

@@ -0,0 +1,60 @@
+import pyttsx3
+import queue
+import time
+import threading
+from ..setting.setting import AppConfig
+
+class Speaker:
+    _instance = None
+    _lock = threading.Lock()  # 确保线程安全的单例模式
+    def __new__(cls, *args, **kwargs):
+        """单例模式：确保全局只有一个配置实例"""
+        if not cls._instance:
+            with cls._lock:
+                if not cls._instance:
+                    cls._instance = super().__new__(cls)
+                    # 标记是否已经初始化，防止 __init__ 重复调用
+                    cls._instance._initialized = False
+        return cls._instance
+    
+    def __init__(self,volum = None ,speed = None) -> None:
+        if self._initialized:
+            return
+        
+        config_manager = AppConfig()
+        self.volum = volum if volum else config_manager.volume
+        self.speed = speed if speed else config_manager.speed
+        
+        self.voice_tasks = queue.Queue()
+        
+        # 启动工作线程
+        self.thread = threading.Thread(target=self.speak_text_block_thread, daemon=True)
+        self.thread.start()
+        
+        self._initialized = True
+        
+    def add_sentence(self, sentence: str):
+        self.voice_tasks.put(sentence)
+    
+    def speak_text_block_thread(self):
+        # 【关键修改】在子线程内部初始化引擎
+        engine = pyttsx3.init()
+        engine.setProperty('rate',self.speed ) 
+        
+        while True:
+            # queue.get() 默认是阻塞的，没有任务时会一直等在这里，不会消耗CPU
+            speak_sentence = self.voice_tasks.get()
+            
+            if speak_sentence:
+                engine.say(speak_sentence)
+                engine.runAndWait()
+            
+            # 通知队列当前任务已处理完毕（良好的队列使用习惯）
+            self.voice_tasks.task_done()
+                
+if __name__ == "__main__":
+    sp = Speaker()
+    for i in range(99):
+        sp.add_sentence(f"阿巴阿巴阿巴 {i}")
+        
+    time.sleep(100)