About Local offline translation program, support text translation, underlining translation, screenshot translation, voice (audio file) translation, video translation, txt file, PPT, Word, PDF, Excel, picture translation.  Resources

from DocImgTranslation.FastSAM.fastsam import FastSAM, FastSAMPrompt
import ast
import torch
from PIL import Image
import configparser
config = configparser.ConfigParser()
config.read("config.conf")
device = config.get("devices", "device")
def main(img_path,device=device,text_prompt="A sheet of white A4 paper",model_path='models/FastSAM.pt'):
    # load model
    model = FastSAM(model_path)
    input = Image.open(img_path)

    # 检查图像的旋转方向标志
    if hasattr(input, '_getexif'):  # 检查是否有EXIF数据
        exif = input._getexif()
        if exif is not None:
            orientation = exif.get(0x0112)  # 获取方向标志（0x0112是EXIF中的方向标志）
            if orientation is not None:
                if orientation == 1:  # 没有旋转
                    pass
                elif orientation == 3:  # 旋转180度
                    input = input.rotate(180, expand=True)
                elif orientation == 6:  # 顺时针旋转90度
                    input = input.rotate(270, expand=True)
                elif orientation == 8:  # 逆时针旋转90度
                    input = input.rotate(90, expand=True)
    input = input.convert("RGB")
    everything_results = model(
        input,
        device=device,
        retina_masks=True,
        imgsz=512,
        conf=0.7,
        iou=0.6
        )
    prompt_process = FastSAMPrompt(input, everything_results, device=device)
    # ann = prompt_process.text_prompt(text=text_prompt)
    # everything prompt
    ann = prompt_process.everything_prompt()
    return ann

if __name__ == "__main__":
    img_path="/Users/liuhongdi/计算机设计大赛/OCR图片文字识别/数据/jpg_pdf/1/16.jpg"
    main(img_path,device='cpu')

DocImgTranslation/FastSAM/Inference.py
tianclll/Ace-Translate
from ultralytics import YOLO
import gradio as gr
import torch
from utils.tools_gradio import fast_process
from utils.tools import format_results, box_prompt, point_prompt, text_prompt
from PIL import ImageDraw
import numpy as np

# Load the pre-trained model
model = YOLO('./weights/FastSAM.pt')

device = torch.device(
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)

# Description
title = "<center><strong><font size='8'>🏃 Fast Segment Anything 🤗</font></strong></center>"

news = """ # 📖 News
        🔥 2023/07/14: Add a "wider result" button in text mode (Thanks for [gaoxinge](https://github.com/CASIA-IVA-Lab/FastSAM/pull/95)).

        🔥 2023/06/29: Support the text mode (Thanks for [gaoxinge](https://github.com/CASIA-IVA-Lab/FastSAM/pull/47)).

        🔥 2023/06/26: Support the points mode. (Better and faster interaction will come soon!)

        🔥 2023/06/24: Add the 'Advanced options" in Everything mode to get a more detailed adjustment.        
        """  

description_e = """This is a demo on Github project 🏃 [Fast Segment Anything Model](https://github.com/CASIA-IVA-Lab/FastSAM). Welcome to give a star ⭐️ to it.
                
                🎯 Upload an Image, segment it with Fast Segment Anything (Everything mode). The other modes will come soon.
                
                ⌛️ It takes about 6~ seconds to generate segment results. The concurrency_count of queue is 1, please wait for a moment when it is crowded.
                
                🚀 To get faster results, you can use a smaller input size and leave high_visual_quality unchecked.
                
                📣 You can also obtain the segmentation results of any Image through this Colab: [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1oX14f6IneGGw612WgVlAiy91UHwFAvr9?usp=sharing)
                
                😚 A huge thanks goes out to the @HuggingFace Team for supporting us with GPU grant.
                
                🏠 Check out our [Model Card 🏃](https://huggingface.co/An-619/FastSAM)
                
              """

description_p = """ # 🎯 Instructions for points mode
                This is a demo on Github project 🏃 [Fast Segment Anything Model](https://github.com/CASIA-IVA-Lab/FastSAM). Welcome to give a star ⭐️ to it.
                
                1. Upload an image or choose an example.
                
                2. Choose the point label ('Add mask' means a positive point. 'Remove' Area means a negative point that is not segmented).
                
                3. Add points one by one on the image.
                
                4. Click the 'Segment with points prompt' button to get the segmentation results.
                
                **5. If you get Error, click the 'Clear points' button and try again may help.**
                
              """

examples = [["examples/sa_8776.jpg"], ["examples/sa_414.jpg"], ["examples/sa_1309.jpg"], ["examples/sa_11025.jpg"],
            ["examples/sa_561.jpg"], ["examples/sa_192.jpg"], ["examples/sa_10039.jpg"], ["examples/sa_862.jpg"]]

default_example = examples[0]

css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"


def segment_everything(
    input,
    input_size=1024, 
    iou_threshold=0.7,
    conf_threshold=0.25,
    better_quality=False,
    withContours=True,
    use_retina=True,
    text="",
    wider=False,
    mask_random_color=True,
):
    input_size = int(input_size)  # 确保 imgsz 是整数
    # Thanks for the suggestion by hysts in HuggingFace.
    w, h = input.size
    scale = input_size / max(w, h)
    new_w = int(w * scale)
    new_h = int(h * scale)
    input = input.resize((new_w, new_h))

    results = model(input,
                    device=device,
                    retina_masks=True,
                    iou=iou_threshold,
                    conf=conf_threshold,
                    imgsz=input_size,)

    if len(text) > 0:
        results = format_results(results[0], 0)
        annotations, _ = text_prompt(results, text, input, device=device, wider=wider)
        annotations = np.array([annotations])
    else:
        annotations = results[0].masks.data
    
    fig = fast_process(annotations=annotations,
                       image=input,
                       device=device,
                       scale=(1024 // input_size),
                       better_quality=better_quality,
                       mask_random_color=mask_random_color,
                       bbox=None,
                       use_retina=use_retina,
                       withContours=withContours,)
    return fig


def segment_with_points(
    input,
    input_size=1024, 
    iou_threshold=0.7,
    conf_threshold=0.25,
    better_quality=False,
    withContours=True,
    use_retina=True,
    mask_random_color=True,
):
    global global_points
    global global_point_label
    
    input_size = int(input_size)  # 确保 imgsz 是整数
    # Thanks for the suggestion by hysts in HuggingFace.
    w, h = input.size
    scale = input_size / max(w, h)
    new_w = int(w * scale)
    new_h = int(h * scale)
    input = input.resize((new_w, new_h))
    
    scaled_points = [[int(x * scale) for x in point] for point in global_points]

    results = model(input,
                    device=device,
                    retina_masks=True,
                    iou=iou_threshold,
                    conf=conf_threshold,
                    imgsz=input_size,)
    
    results = format_results(results[0], 0)
    annotations, _ = point_prompt(results, scaled_points, global_point_label, new_h, new_w)
    annotations = np.array([annotations])

    fig = fast_process(annotations=annotations,
                       image=input,
                       device=device,
                       scale=(1024 // input_size),
                       better_quality=better_quality,
                       mask_random_color=mask_random_color,
                       bbox=None,
                       use_retina=use_retina,
                       withContours=withContours,)

    global_points = []
    global_point_label = []
    return fig, None


def get_points_with_draw(image, label, evt: gr.SelectData):
    global global_points
    global global_point_label

    x, y = evt.index[0], evt.index[1]
    point_radius, point_color = 15, (255, 255, 0) if label == 'Add Mask' else (255, 0, 255)
    global_points.append([x, y])
    global_point_label.append(1 if label == 'Add Mask' else 0)
    
    print(x, y, label == 'Add Mask')
    
    # 创建一个可以在图像上绘图的对象
    draw = ImageDraw.Draw(image)
    draw.ellipse([(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)], fill=point_color)
    return image


cond_img_e = gr.Image(label="Input", value=default_example[0], type='pil')
cond_img_p = gr.Image(label="Input with points", value=default_example[0], type='pil')
cond_img_t = gr.Image(label="Input with text", value="examples/dogs.jpg", type='pil')

segm_img_e = gr.Image(label="Segmented Image", interactive=False, type='pil')
segm_img_p = gr.Image(label="Segmented Image with points", interactive=False, type='pil')
segm_img_t = gr.Image(label="Segmented Image with text", interactive=False, type='pil')

global_points = []
global_point_label = []

input_size_slider = gr.components.Slider(minimum=512,
                                         maximum=1024,
                                         value=1024,
                                         step=64,
                                         label='Input_size',
                                         info='Our model was trained on a size of 1024')

with gr.Blocks(css=css, title='Fast Segment Anything') as demo:
    with gr.Row():
        with gr.Column(scale=1):
            # Title
            gr.Markdown(title)

        with gr.Column(scale=1):
            # News
            gr.Markdown(news)

    with gr.Tab("Everything mode"):
        # Images
        with gr.Row(variant="panel"):
            with gr.Column(scale=1):
                cond_img_e.render()

            with gr.Column(scale=1):
                segm_img_e.render()

        # Submit & Clear
        with gr.Row():
            with gr.Column():
                input_size_slider.render()

                with gr.Row():
                    contour_check = gr.Checkbox(value=True, label='withContours', info='draw the edges of the masks')

                    with gr.Column():
                        segment_btn_e = gr.Button("Segment Everything", variant='primary')
                        clear_btn_e = gr.Button("Clear", variant="secondary")

                gr.Markdown("Try some of the examples below ⬇️")
                gr.Examples(examples=examples,
                            inputs=[cond_img_e],
                            outputs=segm_img_e,
                            fn=segment_everything,
                            cache_examples=True,
                            examples_per_page=4)

            with gr.Column():
                with gr.Accordion("Advanced options", open=False):
                    iou_threshold = gr.Slider(0.1, 0.9, 0.7, step=0.1, label='iou', info='iou threshold for filtering the annotations')
                    conf_threshold = gr.Slider(0.1, 0.9, 0.25, step=0.05, label='conf', info='object confidence threshold')
                    with gr.Row():
                        mor_check = gr.Checkbox(value=False, label='better_visual_quality', info='better quality using morphologyEx')
                        with gr.Column():
                            retina_check = gr.Checkbox(value=True, label='use_retina', info='draw high-resolution segmentation masks')

                # Description
                gr.Markdown(description_e)

    segment_btn_e.click(segment_everything,
                        inputs=[
                            cond_img_e,
                            input_size_slider,
                            iou_threshold,
                            conf_threshold,
                            mor_check,
                            contour_check,
                            retina_check,
                        ],
                        outputs=segm_img_e)

    with gr.Tab("Points mode"):
        # Images
        with gr.Row(variant="panel"):
            with gr.Column(scale=1):
                cond_img_p.render()

            with gr.Column(scale=1):
                segm_img_p.render()
                
        # Submit & Clear
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    add_or_remove = gr.Radio(["Add Mask", "Remove Area"], value="Add Mask", label="Point_label (foreground/background)")

                    with gr.Column():
                        segment_btn_p = gr.Button("Segment with points prompt", variant='primary')
                        clear_btn_p = gr.Button("Clear points", variant='secondary')

                gr.Markdown("Try some of the examples below ⬇️")
                gr.Examples(examples=examples,
                            inputs=[cond_img_p],
                            # outputs=segm_img_p,
                            # fn=segment_with_points,
                            # cache_examples=True,
                            examples_per_page=4)

            with gr.Column():
                # Description
                gr.Markdown(description_p)

    cond_img_p.select(get_points_with_draw, [cond_img_p, add_or_remove], cond_img_p)

    segment_btn_p.click(segment_with_points,
                        inputs=[cond_img_p],
                        outputs=[segm_img_p, cond_img_p])

    with gr.Tab("Text mode"):
        # Images
        with gr.Row(variant="panel"):
            with gr.Column(scale=1):
                cond_img_t.render()

            with gr.Column(scale=1):
                segm_img_t.render()

        # Submit & Clear
        with gr.Row():
            with gr.Column():
                input_size_slider_t = gr.components.Slider(minimum=512,
                                                           maximum=1024,
                                                           value=1024,
                                                           step=64,
                                                           label='Input_size',
                                                           info='Our model was trained on a size of 1024')
                with gr.Row():
                    with gr.Column():
                        contour_check = gr.Checkbox(value=True, label='withContours', info='draw the edges of the masks')
                        text_box = gr.Textbox(label="text prompt", value="a black dog")

                    with gr.Column():
                        segment_btn_t = gr.Button("Segment with text", variant='primary')
                        clear_btn_t = gr.Button("Clear", variant="secondary")

                gr.Markdown("Try some of the examples below ⬇️")
                gr.Examples(examples=[["examples/dogs.jpg"]] + examples,
                            inputs=[cond_img_e],
                            # outputs=segm_img_e,
                            # fn=segment_everything,
                            # cache_examples=True,
                            examples_per_page=4)

            with gr.Column():
                with gr.Accordion("Advanced options", open=False):
                    iou_threshold = gr.Slider(0.1, 0.9, 0.7, step=0.1, label='iou', info='iou threshold for filtering the annotations')
                    conf_threshold = gr.Slider(0.1, 0.9, 0.25, step=0.05, label='conf', info='object confidence threshold')
                    with gr.Row():
                        mor_check = gr.Checkbox(value=False, label='better_visual_quality', info='better quality using morphologyEx')
                        retina_check = gr.Checkbox(value=True, label='use_retina', info='draw high-resolution segmentation masks')
                        wider_check = gr.Checkbox(value=False, label='wider', info='wider result')

                # Description
                gr.Markdown(description_e)
    
    segment_btn_t.click(segment_everything,
                        inputs=[
                            cond_img_t,
                            input_size_slider_t,
                            iou_threshold,
                            conf_threshold,
                            mor_check,
                            contour_check,
                            retina_check,
                            text_box,
                            wider_check,
                        ],
                        outputs=segm_img_t)

    def clear():
        return None, None
    
    def clear_text():
        return None, None, None

    clear_btn_e.click(clear, outputs=[cond_img_e, segm_img_e])
    clear_btn_p.click(clear, outputs=[cond_img_p, segm_img_p])
    clear_btn_t.click(clear_text, outputs=[cond_img_p, segm_img_p, text_box])

demo.queue()
demo.launch()

DocImgTranslation/FastSAM/app_gradio.py
tianclll/Ace-Translate
from .model import FastSAM
import numpy as np
from PIL import Image
import clip
from typing import Optional, List, Tuple, Union


class FastSAMDecoder:
    def __init__(
        self,
        model: FastSAM,
        device: str='cpu',
        conf: float=0.4, 
        iou: float=0.9,
        imgsz: int=1024,
        retina_masks: bool=True,
        ):
        self.model = model
        self.device = device
        self.retina_masks = retina_masks
        self.imgsz = imgsz
        self.conf = conf
        self.iou = iou
        self.image = None
        self.image_embedding = None
        
    def run_encoder(self, image):
        if isinstance(image,str):
            image =  np.array(Image.open(image))
        self.image = image
        image_embedding = self.model(
            self.image, 
            device=self.device, 
            retina_masks=self.retina_masks, 
            imgsz=self.imgsz, 
            conf=self.conf, 
            iou=self.iou
            )
        return image_embedding[0].numpy()

    def run_decoder(
            self,
            image_embedding, 
            point_prompt: Optional[np.ndarray]=None,
            point_label: Optional[np.ndarray]=None,
            box_prompt: Optional[np.ndarray]=None,
            text_prompt: Optional[str]=None,
            )->np.ndarray:
        self.image_embedding = image_embedding
        if point_prompt is not None:
            ann = self.point_prompt(points=point_prompt, pointlabel=point_label)
            return ann
        elif box_prompt is not None:
            ann = self.box_prompt(bbox=box_prompt)
            return ann
        elif text_prompt is not None:
            ann = self.text_prompt(text=text_prompt)
            return ann
        else:
            return None

    def box_prompt(self, bbox):
        assert (bbox[2] != 0 and bbox[3] != 0)
        masks = self.image_embedding.masks.data
        target_height = self.image.shape[0]
        target_width = self.image.shape[1]
        h = masks.shape[1]
        w = masks.shape[2]
        if h != target_height or w != target_width:
            bbox = [
                int(bbox[0] * w / target_width),
                int(bbox[1] * h / target_height),
                int(bbox[2] * w / target_width),
                int(bbox[3] * h / target_height), ]
        bbox[0] = round(bbox[0]) if round(bbox[0]) > 0 else 0
        bbox[1] = round(bbox[1]) if round(bbox[1]) > 0 else 0
        bbox[2] = round(bbox[2]) if round(bbox[2]) < w else w
        bbox[3] = round(bbox[3]) if round(bbox[3]) < h else h

        # IoUs = torch.zeros(len(masks), dtype=torch.float32)
        bbox_area = (bbox[3] - bbox[1]) * (bbox[2] - bbox[0])

        masks_area = np.sum(masks[:, bbox[1]:bbox[3], bbox[0]:bbox[2]], axis=(1, 2))
        orig_masks_area = np.sum(masks, axis=(1, 2))

        union = bbox_area + orig_masks_area - masks_area
        IoUs = masks_area / union
        max_iou_index = np.argmax(IoUs)

        return np.array([masks[max_iou_index].cpu().numpy()])

    def point_prompt(self, points, pointlabel):  # numpy 

        masks = self._format_results(self.image_embedding[0], 0)
        target_height = self.image.shape[0]
        target_width = self.image.shape[1]
        h = masks[0]['segmentation'].shape[0]
        w = masks[0]['segmentation'].shape[1]
        if h != target_height or w != target_width:
            points = [[int(point[0] * w / target_width), int(point[1] * h / target_height)] for point in points]
        onemask = np.zeros((h, w))
        masks = sorted(masks, key=lambda x: x['area'], reverse=True)
        for i, annotation in enumerate(masks):
            if type(annotation) == dict:
                mask = annotation['segmentation']
            else:
                mask = annotation
            for i, point in enumerate(points):
                if mask[point[1], point[0]] == 1 and pointlabel[i] == 1:
                    onemask[mask] = 1
                if mask[point[1], point[0]] == 1 and pointlabel[i] == 0:
                    onemask[mask] = 0
        onemask = onemask >= 1
        return np.array([onemask])
    
    def _format_results(self, result, filter=0):
        annotations = []
        n = len(result.masks.data)
        for i in range(n):
            annotation = {}
            mask = result.masks.data[i] == 1.0

            if np.sum(mask) < filter:
                continue
            annotation['id'] = i
            annotation['segmentation'] = mask
            annotation['bbox'] = result.boxes.data[i]
            annotation['score'] = result.boxes.conf[i]
            annotation['area'] = annotation['segmentation'].sum()
            annotations.append(annotation)
        return annotations

DocImgTranslation/FastSAM/fastsam/decoder.py
tianclll/Ace-Translate
# Ultralytics YOLO 🚀, AGPL-3.0 license
"""
FastSAM model interface.

Usage - Predict:
    from ultralytics import FastSAM

    model = FastSAM('last.pt')
    results = model.predict('ultralytics/assets/bus.jpg')
"""

from ultralytics.yolo.cfg import get_cfg
from ultralytics.yolo.engine.exporter import Exporter
from ultralytics.yolo.engine.model import YOLO
from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, ROOT, is_git_dir
from ultralytics.yolo.utils.checks import check_imgsz

from ultralytics.yolo.utils.torch_utils import model_info, smart_inference_mode
from .predict import FastSAMPredictor


class FastSAM(YOLO):

    @smart_inference_mode()
    def predict(self, source=None, stream=False, **kwargs):
        """
        Perform prediction using the YOLO model.

        Args:
            source (str | int | PIL | np.ndarray): The source of the image to make predictions on.
                          Accepts all source types accepted by the YOLO model.
            stream (bool): Whether to stream the predictions or not. Defaults to False.
            **kwargs : Additional keyword arguments passed to the predictor.
                       Check the 'configuration' section in the documentation for all available options.

        Returns:
            (List[ultralytics.yolo.engine.results.Results]): The prediction results.
        """
        if source is None:
            source = ROOT / 'assets' if is_git_dir() else 'https://ultralytics.com/images/bus.jpg'
            LOGGER.warning(f"WARNING ⚠️ 'source' is missing. Using 'source={source}'.")
        overrides = self.overrides.copy()
        overrides['conf'] = 0.25
        overrides.update(kwargs)  # prefer kwargs
        overrides['mode'] = kwargs.get('mode', 'predict')
        assert overrides['mode'] in ['track', 'predict']
        overrides['save'] = kwargs.get('save', False)  # do not save by default if called in Python
        self.predictor = FastSAMPredictor(overrides=overrides)
        self.predictor.setup_model(model=self.model, verbose=False)
        try:
            return self.predictor(source, stream=stream)
        except Exception as e:
            return None

    def train(self, **kwargs):
        """Function trains models but raises an error as FastSAM models do not support training."""
        raise NotImplementedError("Currently, the training codes are on the way.")

    def val(self, **kwargs):
        """Run validation given dataset."""
        overrides = dict(task='segment', mode='val')
        overrides.update(kwargs)  # prefer kwargs
        args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides)
        args.imgsz = check_imgsz(args.imgsz, max_dim=1)
        validator = FastSAM(args=args)
        validator(model=self.model)
        self.metrics = validator.metrics
        return validator.metrics

    @smart_inference_mode()
    def export(self, **kwargs):
        """
        Export model.

        Args:
            **kwargs : Any other args accepted by the predictors. To see all args check 'configuration' section in docs
        """
        overrides = dict(task='detect')
        overrides.update(kwargs)
        overrides['mode'] = 'export'
        args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides)
        args.task = self.task
        if args.imgsz == DEFAULT_CFG.imgsz:
            args.imgsz = self.model.args['imgsz']  # use trained imgsz unless custom value is passed
        if args.batch == DEFAULT_CFG.batch:
            args.batch = 1  # default to 1 if not modified
        return Exporter(overrides=args)(model=self.model)

    def info(self, detailed=False, verbose=True):
        """
        Logs model info.

        Args:
            detailed (bool): Show detailed information about model.
            verbose (bool): Controls verbosity.
        """
        return model_info(self.model, detailed=detailed, verbose=verbose, imgsz=640)

    def __call__(self, source=None, stream=False, **kwargs):
        """Calls the 'predict' function with given arguments to perform object detection."""
        return self.predict(source, stream, **kwargs)

    def __getattr__(self, attr):
        """Raises error if object has no requested attribute."""
        name = self.__class__.__name__
        raise AttributeError(f"'{name}' object has no attribute '{attr}'. See valid attributes below.\n{self.__doc__}")

DocImgTranslation/FastSAM/fastsam/model.py
tianclll/Ace-Translate
import torch

from ultralytics.yolo.engine.results import Results
from ultralytics.yolo.utils import DEFAULT_CFG, ops
from ultralytics.yolo.v8.detect.predict import DetectionPredictor
from .utils import bbox_iou

class FastSAMPredictor(DetectionPredictor):

    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
        super().__init__(cfg, overrides, _callbacks)
        self.args.task = 'segment'

    def postprocess(self, preds, img, orig_imgs):
        """TODO: filter by classes."""
        p = ops.non_max_suppression(preds[0],
                                    self.args.conf,
                                    self.args.iou,
                                    agnostic=self.args.agnostic_nms,
                                    max_det=self.args.max_det,
                                    nc=len(self.model.names),
                                    classes=self.args.classes)

        results = []
        if len(p) == 0 or len(p[0]) == 0:
            print("No object detected.")
            return results

        full_box = torch.zeros_like(p[0][0])
        full_box[2], full_box[3], full_box[4], full_box[6:] = img.shape[3], img.shape[2], 1.0, 1.0
        full_box = full_box.view(1, -1)
        critical_iou_index = bbox_iou(full_box[0][:4], p[0][:, :4], iou_thres=0.9, image_shape=img.shape[2:])
        if critical_iou_index.numel() != 0:
            full_box[0][4] = p[0][critical_iou_index][:,4]
            full_box[0][6:] = p[0][critical_iou_index][:,6:]
            p[0][critical_iou_index] = full_box
        
        proto = preds[1][-1] if len(preds[1]) == 3 else preds[1]  # second output is len 3 if pt, but only 1 if exported
        for i, pred in enumerate(p):
            orig_img = orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs
            path = self.batch[0]
            img_path = path[i] if isinstance(path, list) else path
            if not len(pred):  # save empty boxes
                results.append(Results(orig_img=orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6]))
                continue
            if self.args.retina_masks:
                if not isinstance(orig_imgs, torch.Tensor):
                    pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
                masks = ops.process_mask_native(proto[i], pred[:, 6:], pred[:, :4], orig_img.shape[:2])  # HWC
            else:
                masks = ops.process_mask(proto[i], pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True)  # HWC
                if not isinstance(orig_imgs, torch.Tensor):
                    pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
            results.append(
                Results(orig_img=orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks))
        return results

DocImgTranslation/FastSAM/fastsam/predict.py
tianclll/Ace-Translate
import os
import sys
import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
from .utils import image_to_np_ndarray
from PIL import Image

try:
    import clip  # for linear_assignment

except (ImportError, AssertionError, AttributeError):
    from ultralytics.yolo.utils.checks import check_requirements

    check_requirements('git+https://github.com/openai/CLIP.git')  # required before installing lap from source
    import clip


class FastSAMPrompt:

    def __init__(self, image, results, device='cuda'):
        if isinstance(image, str) or isinstance(image, Image.Image):
            image = image_to_np_ndarray(image)
        self.device = device
        self.results = results
        self.img = image
    
    def _segment_image(self, image, bbox):
        if isinstance(image, Image.Image):
            image_array = np.array(image)
        else:
            image_array = image
        segmented_image_array = np.zeros_like(image_array)
        x1, y1, x2, y2 = bbox
        segmented_image_array[y1:y2, x1:x2] = image_array[y1:y2, x1:x2]
        segmented_image = Image.fromarray(segmented_image_array)
        black_image = Image.new('RGB', image.size, (255, 255, 255))
        # transparency_mask = np.zeros_like((), dtype=np.uint8)
        transparency_mask = np.zeros((image_array.shape[0], image_array.shape[1]), dtype=np.uint8)
        transparency_mask[y1:y2, x1:x2] = 255
        transparency_mask_image = Image.fromarray(transparency_mask, mode='L')
        black_image.paste(segmented_image, mask=transparency_mask_image)
        return black_image

    def _format_results(self, result, filter=0):
        annotations = []
        n = len(result.masks.data)
        for i in range(n):
            annotation = {}
            mask = result.masks.data[i] == 1.0

            if torch.sum(mask) < filter:
                continue
            annotation['id'] = i
            annotation['segmentation'] = mask.cpu().numpy()
            annotation['bbox'] = result.boxes.data[i]
            annotation['score'] = result.boxes.conf[i]
            annotation['area'] = annotation['segmentation'].sum()
            annotations.append(annotation)
        return annotations

    def filter_masks(annotations):  # filte the overlap mask
        annotations.sort(key=lambda x: x['area'], reverse=True)
        to_remove = set()
        for i in range(0, len(annotations)):
            a = annotations[i]
            for j in range(i + 1, len(annotations)):
                b = annotations[j]
                if i != j and j not in to_remove:
                    # check if
                    if b['area'] < a['area']:
                        if (a['segmentation'] & b['segmentation']).sum() / b['segmentation'].sum() > 0.8:
                            to_remove.add(j)

        return [a for i, a in enumerate(annotations) if i not in to_remove], to_remove

    def _get_bbox_from_mask(self, mask):
        mask = mask.astype(np.uint8)
        contours, hierarchy = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        x1, y1, w, h = cv2.boundingRect(contours[0])
        x2, y2 = x1 + w, y1 + h
        if len(contours) > 1:
            for b in contours:
                x_t, y_t, w_t, h_t = cv2.boundingRect(b)
                # Merge multiple bounding boxes into one.
                x1 = min(x1, x_t)
                y1 = min(y1, y_t)
                x2 = max(x2, x_t + w_t)
                y2 = max(y2, y_t + h_t)
            h = y2 - y1
            w = x2 - x1
        return [x1, y1, x2, y2]

    def plot_to_result(self,
             annotations,
             bboxes=None,
             points=None,
             point_label=None,
             mask_random_color=True,
             better_quality=True,
             retina=False,
             withContours=True) -> np.ndarray:
        if isinstance(annotations[0], dict):
            annotations = [annotation['segmentation'] for annotation in annotations]
        image = self.img
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        original_h = image.shape[0]
        original_w = image.shape[1]
        if sys.platform == "darwin":
            plt.switch_backend("TkAgg")
        plt.figure(figsize=(original_w / 100, original_h / 100))
        # Add subplot with no margin.
        plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
        plt.margins(0, 0)
        plt.gca().xaxis.set_major_locator(plt.NullLocator())
        plt.gca().yaxis.set_major_locator(plt.NullLocator())

        plt.imshow(image)
        if better_quality:
            if isinstance(annotations[0], torch.Tensor):
                annotations = np.array(annotations.cpu())
            for i, mask in enumerate(annotations):
                mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8))
                annotations[i] = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8))
        if self.device == 'cpu':
            annotations = np.array(annotations)
            self.fast_show_mask(
                annotations,
                plt.gca(),
                random_color=mask_random_color,
                bboxes=bboxes,
                points=points,
                pointlabel=point_label,
                retinamask=retina,
                target_height=original_h,
                target_width=original_w,
            )
        else:
            if isinstance(annotations[0], np.ndarray):
                annotations = torch.from_numpy(annotations)
            self.fast_show_mask_gpu(
                annotations,
                plt.gca(),
                random_color=mask_random_color,
                bboxes=bboxes,
                points=points,
                pointlabel=point_label,
                retinamask=retina,
                target_height=original_h,
                target_width=original_w,
            )
        if isinstance(annotations, torch.Tensor):
            annotations = annotations.cpu().numpy()
        if withContours:
            contour_all = []
            temp = np.zeros((original_h, original_w, 1))
            for i, mask in enumerate(annotations):
                if type(mask) == dict:
                    mask = mask['segmentation']
                annotation = mask.astype(np.uint8)
                if not retina:
                    annotation = cv2.resize(
                        annotation,
                        (original_w, original_h),
                        interpolation=cv2.INTER_NEAREST,
                    )
                contours, hierarchy = cv2.findContours(annotation, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
                for contour in contours:
                    contour_all.append(contour)
            cv2.drawContours(temp, contour_all, -1, (255, 255, 255), 2)
            color = np.array([0 / 255, 0 / 255, 255 / 255, 0.8])
            contour_mask = temp / 255 * color.reshape(1, 1, -1)
            plt.imshow(contour_mask)

        plt.axis('off')
        fig = plt.gcf()
        plt.draw()

        try:
            buf = fig.canvas.tostring_rgb()
        except AttributeError:
            fig.canvas.draw()
            buf = fig.canvas.tostring_rgb()
        cols, rows = fig.canvas.get_width_height()
        img_array = np.frombuffer(buf, dtype=np.uint8).reshape(rows, cols, 3)
        result = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
        plt.close()
        return result
            
    # Remark for refactoring: IMO a function should do one thing only, storing the image and plotting should be seperated and do not necessarily need to be class functions but standalone utility functions that the user can chain in his scripts to have more fine-grained control. 
    def plot(self,
             annotations,
             output_path,
             bboxes=None,
             points=None,
             point_label=None,
             mask_random_color=True,
             better_quality=True,
             retina=False,
             withContours=True):
        if len(annotations) == 0:
            return None
        result = self.plot_to_result(
            annotations, 
            bboxes, 
            points, 
            point_label, 
            mask_random_color,
            better_quality, 
            retina, 
            withContours,
        )

        path = os.path.dirname(os.path.abspath(output_path))
        if not os.path.exists(path):
            os.makedirs(path)
        result = result[:, :, ::-1]
        cv2.imwrite(output_path, result)
     
    #   CPU post process
    def fast_show_mask(
        self,
        annotation,
        ax,
        random_color=False,
        bboxes=None,
        points=None,
        pointlabel=None,
        retinamask=True,
        target_height=960,
        target_width=960,
    ):
        msak_sum = annotation.shape[0]
        height = annotation.shape[1]
        weight = annotation.shape[2]
        #Sort annotations based on area.
        areas = np.sum(annotation, axis=(1, 2))
        sorted_indices = np.argsort(areas)
        annotation = annotation[sorted_indices]

        index = (annotation != 0).argmax(axis=0)
        if random_color:
            color = np.random.random((msak_sum, 1, 1, 3))
        else:
            color = np.ones((msak_sum, 1, 1, 3)) * np.array([30 / 255, 144 / 255, 255 / 255])
        transparency = np.ones((msak_sum, 1, 1, 1)) * 0.6
        visual = np.concatenate([color, transparency], axis=-1)
        mask_image = np.expand_dims(annotation, -1) * visual

        show = np.zeros((height, weight, 4))
        h_indices, w_indices = np.meshgrid(np.arange(height), np.arange(weight), indexing='ij')
        indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
        # Use vectorized indexing to update the values of 'show'.
        show[h_indices, w_indices, :] = mask_image[indices]
        if bboxes is not None:
            for bbox in bboxes:
                x1, y1, x2, y2 = bbox
                ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor='b', linewidth=1))
        # draw point
        if points is not None:
            plt.scatter(
                [point[0] for i, point in enumerate(points) if pointlabel[i] == 1],
                [point[1] for i, point in enumerate(points) if pointlabel[i] == 1],
                s=20,
                c='y',
            )
            plt.scatter(
                [point[0] for i, point in enumerate(points) if pointlabel[i] == 0],
                [point[1] for i, point in enumerate(points) if pointlabel[i] == 0],
                s=20,
                c='m',
            )

        if not retinamask:
            show = cv2.resize(show, (target_width, target_height), interpolation=cv2.INTER_NEAREST)
        ax.imshow(show)

    def fast_show_mask_gpu(
        self,
        annotation,
        ax,
        random_color=False,
        bboxes=None,
        points=None,
        pointlabel=None,
        retinamask=True,
        target_height=960,
        target_width=960,
    ):
        msak_sum = annotation.shape[0]
        height = annotation.shape[1]
        weight = annotation.shape[2]
        areas = torch.sum(annotation, dim=(1, 2))
        sorted_indices = torch.argsort(areas, descending=False)
        annotation = annotation[sorted_indices]
        # Find the index of the first non-zero value at each position.
        index = (annotation != 0).to(torch.long).argmax(dim=0)
        if random_color:
            color = torch.rand((msak_sum, 1, 1, 3)).to(annotation.device)
        else:
            color = torch.ones((msak_sum, 1, 1, 3)).to(annotation.device) * torch.tensor([
                30 / 255, 144 / 255, 255 / 255]).to(annotation.device)
        transparency = torch.ones((msak_sum, 1, 1, 1)).to(annotation.device) * 0.6
        visual = torch.cat([color, transparency], dim=-1)
        mask_image = torch.unsqueeze(annotation, -1) * visual
        # Select data according to the index. The index indicates which batch's data to choose at each position, converting the mask_image into a single batch form.
        show = torch.zeros((height, weight, 4)).to(annotation.device)
        try:
            h_indices, w_indices = torch.meshgrid(torch.arange(height), torch.arange(weight), indexing='ij')
        except:
            h_indices, w_indices = torch.meshgrid(torch.arange(height), torch.arange(weight))
        indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
        # Use vectorized indexing to update the values of 'show'.
        show[h_indices, w_indices, :] = mask_image[indices]
        show_cpu = show.cpu().numpy()
        if bboxes is not None:
            for bbox in bboxes:
                x1, y1, x2, y2 = bbox
                ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor='b', linewidth=1))
        # draw point
        if points is not None:
            plt.scatter(
                [point[0] for i, point in enumerate(points) if pointlabel[i] == 1],
                [point[1] for i, point in enumerate(points) if pointlabel[i] == 1],
                s=20,
                c='y',
            )
            plt.scatter(
                [point[0] for i, point in enumerate(points) if pointlabel[i] == 0],
                [point[1] for i, point in enumerate(points) if pointlabel[i] == 0],
                s=20,
                c='m',
            )
        if not retinamask:
            show_cpu = cv2.resize(show_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST)
        ax.imshow(show_cpu)

    # clip
    @torch.no_grad()
    def retrieve(self, model, preprocess, elements, search_text: str, device) -> int:
        preprocessed_images = [preprocess(image).to(device) for image in elements]
        tokenized_text = clip.tokenize([search_text]).to(device)
        stacked_images = torch.stack(preprocessed_images)
        image_features = model.encode_image(stacked_images)
        text_features = model.encode_text(tokenized_text)
        image_features /= image_features.norm(dim=-1, keepdim=True)
        text_features /= text_features.norm(dim=-1, keepdim=True)
        probs = 100.0 * image_features @ text_features.T
        return probs[:, 0].softmax(dim=0)

    def _crop_image(self, format_results):

        image = Image.fromarray(cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB))
        ori_w, ori_h = image.size
        annotations = format_results
        mask_h, mask_w = annotations[0]['segmentation'].shape
        if ori_w != mask_w or ori_h != mask_h:
            image = image.resize((mask_w, mask_h))
        cropped_boxes = []
        cropped_images = []
        not_crop = []
        filter_id = []
        # annotations, _ = filter_masks(annotations)
        # filter_id = list(_)
        for _, mask in enumerate(annotations):
            if np.sum(mask['segmentation']) <= 100:
                filter_id.append(_)
                continue
            bbox = self._get_bbox_from_mask(mask['segmentation'])  # mask 的 bbox
            cropped_boxes.append(self._segment_image(image, bbox))  
            # cropped_boxes.append(segment_image(image,mask["segmentation"]))
            cropped_images.append(bbox)  # Save the bounding box of the cropped image.

        return cropped_boxes, cropped_images, not_crop, filter_id, annotations

    def box_prompt(self, bbox=None, bboxes=None):
        if self.results == None:
            return []
        assert bbox or bboxes
        if bboxes is None:
            bboxes = [bbox]
        max_iou_index = []
        for bbox in bboxes:
            assert (bbox[2] != 0 and bbox[3] != 0)
            masks = self.results[0].masks.data
            target_height = self.img.shape[0]
            target_width = self.img.shape[1]
            h = masks.shape[1]
            w = masks.shape[2]
            if h != target_height or w != target_width:
                bbox = [
                    int(bbox[0] * w / target_width),
                    int(bbox[1] * h / target_height),
                    int(bbox[2] * w / target_width),
                    int(bbox[3] * h / target_height), ]
            bbox[0] = round(bbox[0]) if round(bbox[0]) > 0 else 0
            bbox[1] = round(bbox[1]) if round(bbox[1]) > 0 else 0
            bbox[2] = round(bbox[2]) if round(bbox[2]) < w else w
            bbox[3] = round(bbox[3]) if round(bbox[3]) < h else h

            # IoUs = torch.zeros(len(masks), dtype=torch.float32)
            bbox_area = (bbox[3] - bbox[1]) * (bbox[2] - bbox[0])

            masks_area = torch.sum(masks[:, bbox[1]:bbox[3], bbox[0]:bbox[2]], dim=(1, 2))
            orig_masks_area = torch.sum(masks, dim=(1, 2))

            union = bbox_area + orig_masks_area - masks_area
            IoUs = masks_area / union
            max_iou_index.append(int(torch.argmax(IoUs)))
        max_iou_index = list(set(max_iou_index))
        return np.array(masks[max_iou_index].cpu().numpy())

    def point_prompt(self, points, pointlabel):  # numpy 
        if self.results == None:
            return []
        masks = self._format_results(self.results[0], 0)
        target_height = self.img.shape[0]
        target_width = self.img.shape[1]
        h = masks[0]['segmentation'].shape[0]
        w = masks[0]['segmentation'].shape[1]
        if h != target_height or w != target_width:
            points = [[int(point[0] * w / target_width), int(point[1] * h / target_height)] for point in points]
        onemask = np.zeros((h, w))
        masks = sorted(masks, key=lambda x: x['area'], reverse=True)
        for i, annotation in enumerate(masks):
            if type(annotation) == dict:
                mask = annotation['segmentation']
            else:
                mask = annotation
            for i, point in enumerate(points):
                if mask[point[1], point[0]] == 1 and pointlabel[i] == 1:
                    onemask[mask] = 1
                if mask[point[1], point[0]] == 1 and pointlabel[i] == 0:
                    onemask[mask] = 0
        onemask = onemask >= 1
        return np.array([onemask])

    def text_prompt(self, text):
        if self.results == None:
            return []
        format_results = self._format_results(self.results[0], 0)
        cropped_boxes, cropped_images, not_crop, filter_id, annotations = self._crop_image(format_results)
        clip_model, preprocess = clip.load('ViT-B/32', device=self.device)
        scores = self.retrieve(clip_model, preprocess, cropped_boxes, text, device=self.device)
        max_idx = scores.argsort()
        max_idx = max_idx[-1]
        max_idx += sum(np.array(filter_id) <= int(max_idx))
        return np.array([annotations[max_idx]['segmentation']])

    def everything_prompt(self):
        if self.results == None:
            return []
        return self.results[0].masks.data
        

DocImgTranslation/FastSAM/fastsam/prompt.py
tianclll/Ace-Translate
import numpy as np
import torch
from PIL import Image


def adjust_bboxes_to_image_border(boxes, image_shape, threshold=20):
    '''Adjust bounding boxes to stick to image border if they are within a certain threshold.
    Args:
    boxes: (n, 4)
    image_shape: (height, width)
    threshold: pixel threshold
    Returns:
    adjusted_boxes: adjusted bounding boxes
    '''

    # Image dimensions
    h, w = image_shape

    # Adjust boxes
    boxes[:, 0] = torch.where(boxes[:, 0] < threshold, torch.tensor(
        0, dtype=torch.float, device=boxes.device), boxes[:, 0])  # x1
    boxes[:, 1] = torch.where(boxes[:, 1] < threshold, torch.tensor(
        0, dtype=torch.float, device=boxes.device), boxes[:, 1])  # y1
    boxes[:, 2] = torch.where(boxes[:, 2] > w - threshold, torch.tensor(
        w, dtype=torch.float, device=boxes.device), boxes[:, 2])  # x2
    boxes[:, 3] = torch.where(boxes[:, 3] > h - threshold, torch.tensor(
        h, dtype=torch.float, device=boxes.device), boxes[:, 3])  # y2

    return boxes



def convert_box_xywh_to_xyxy(box):
    x1 = box[0]
    y1 = box[1]
    x2 = box[0] + box[2]
    y2 = box[1] + box[3]
    return [x1, y1, x2, y2]


def bbox_iou(box1, boxes, iou_thres=0.9, image_shape=(640, 640), raw_output=False):
    '''Compute the Intersection-Over-Union of a bounding box with respect to an array of other bounding boxes.
    Args:
    box1: (4, )
    boxes: (n, 4)
    Returns:
    high_iou_indices: Indices of boxes with IoU > thres
    '''
    boxes = adjust_bboxes_to_image_border(boxes, image_shape)
    # obtain coordinates for intersections
    x1 = torch.max(box1[0], boxes[:, 0])
    y1 = torch.max(box1[1], boxes[:, 1])
    x2 = torch.min(box1[2], boxes[:, 2])
    y2 = torch.min(box1[3], boxes[:, 3])

    # compute the area of intersection
    intersection = (x2 - x1).clamp(0) * (y2 - y1).clamp(0)

    # compute the area of both individual boxes
    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
    box2_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])

    # compute the area of union
    union = box1_area + box2_area - intersection

    # compute the IoU
    iou = intersection / union  # Should be shape (n, )
    if raw_output:
        if iou.numel() == 0:
            return 0
        return iou

    # get indices of boxes with IoU > thres
    high_iou_indices = torch.nonzero(iou > iou_thres).flatten()

    return high_iou_indices


def image_to_np_ndarray(image):
    if type(image) is str:
        return np.array(Image.open(image))
    elif issubclass(type(image), Image.Image):
        return np.array(image)
    elif type(image) is np.ndarray:
        return image
    return None

DocImgTranslation/FastSAM/fastsam/utils.py
tianclll/Ace-Translate
from fastsam import FastSAM, FastSAMPrompt
import torch 

model = FastSAM('FastSAM.pt')
IMAGE_PATH = './images/dogs.jpg'
DEVICE = torch.device(
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
everything_results = model(
    IMAGE_PATH,
    device=DEVICE,
    retina_masks=True,
    imgsz=1024,
    conf=0.4,
    iou=0.9,
)
prompt_process = FastSAMPrompt(IMAGE_PATH, everything_results, device=DEVICE)

# # everything prompt
ann = prompt_process.everything_prompt()

# # bbox prompt
# # bbox default shape [0,0,0,0] -> [x1,y1,x2,y2]
# bboxes default shape [[0,0,0,0]] -> [[x1,y1,x2,y2]]
# ann = prompt_process.box_prompt(bbox=[200, 200, 300, 300])
# ann = prompt_process.box_prompt(bboxes=[[200, 200, 300, 300], [500, 500, 600, 600]])

# # text prompt
# ann = prompt_process.text_prompt(text='a photo of a dog')

# # point prompt
# # points default [[0,0]] [[x1,y1],[x2,y2]]
# # point_label default [0] [1,0] 0:background, 1:foreground
# ann = prompt_process.point_prompt(points=[[620, 360]], pointlabel=[1])

# point prompt
# points default [[0,0]] [[x1,y1],[x2,y2]]
# point_label default [0] [1,0] 0:background, 1:foreground
ann = prompt_process.point_prompt(points=[[620, 360]], pointlabel=[1])

prompt_process.plot(
    annotations=ann,
    output='./output/',
    mask_random_color=True,
    better_quality=True,
    retina=False,
    withContours=True,
)

DocImgTranslation/FastSAM/segpredict.py
tianclll/Ace-Translate
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from setuptools import find_packages, setup

setup(
    name="fastsam",
    version="0.1.1",
    install_requires=[],
    package_dir= {
        "fastsam": "fastsam",
        "fastsam_tools": "utils",
    },
    url="https://github.com/CASIA-IVA-Lab/FastSAM",
)
DocImgTranslation/FastSAM/setup.py
tianclll/Ace-Translate
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import cv2
import torch
import os
import sys
import clip


def convert_box_xywh_to_xyxy(box):
    if len(box) == 4:
        return [box[0], box[1], box[0] + box[2], box[1] + box[3]]
    else:
        result = []
        for b in box:
            b = convert_box_xywh_to_xyxy(b)
            result.append(b)               
    return result


def segment_image(image, bbox):
    image_array = np.array(image)
    segmented_image_array = np.zeros_like(image_array)
    x1, y1, x2, y2 = bbox
    segmented_image_array[y1:y2, x1:x2] = image_array[y1:y2, x1:x2]
    segmented_image = Image.fromarray(segmented_image_array)
    black_image = Image.new("RGB", image.size, (255, 255, 255))
    # transparency_mask = np.zeros_like((), dtype=np.uint8)
    transparency_mask = np.zeros(
        (image_array.shape[0], image_array.shape[1]), dtype=np.uint8
    )
    transparency_mask[y1:y2, x1:x2] = 255
    transparency_mask_image = Image.fromarray(transparency_mask, mode="L")
    black_image.paste(segmented_image, mask=transparency_mask_image)
    return black_image


def format_results(result, filter=0):
    annotations = []
    n = len(result.masks.data)
    for i in range(n):
        annotation = {}
        mask = result.masks.data[i] == 1.0

        if torch.sum(mask) < filter:
            continue
        annotation["id"] = i
        annotation["segmentation"] = mask.cpu().numpy()
        annotation["bbox"] = result.boxes.data[i]
        annotation["score"] = result.boxes.conf[i]
        annotation["area"] = annotation["segmentation"].sum()
        annotations.append(annotation)
    return annotations


def filter_masks(annotations):  # filter the overlap mask
    annotations.sort(key=lambda x: x["area"], reverse=True)
    to_remove = set()
    for i in range(0, len(annotations)):
        a = annotations[i]
        for j in range(i + 1, len(annotations)):
            b = annotations[j]
            if i != j and j not in to_remove:
                # check if
                if b["area"] < a["area"]:
                    if (a["segmentation"] & b["segmentation"]).sum() / b[
                        "segmentation"
                    ].sum() > 0.8:
                        to_remove.add(j)

    return [a for i, a in enumerate(annotations) if i not in to_remove], to_remove


def get_bbox_from_mask(mask):
    mask = mask.astype(np.uint8)
    contours, hierarchy = cv2.findContours(
        mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )
    x1, y1, w, h = cv2.boundingRect(contours[0])
    x2, y2 = x1 + w, y1 + h
    if len(contours) > 1:
        for b in contours:
            x_t, y_t, w_t, h_t = cv2.boundingRect(b)
            # 将多个bbox合并成一个
            x1 = min(x1, x_t)
            y1 = min(y1, y_t)
            x2 = max(x2, x_t + w_t)
            y2 = max(y2, y_t + h_t)
        h = y2 - y1
        w = x2 - x1
    return [x1, y1, x2, y2]


def fast_process(
    annotations, args, mask_random_color, bbox=None, points=None, edges=False
):
    if isinstance(annotations[0], dict):
        annotations = [annotation["segmentation"] for annotation in annotations]
    result_name = os.path.basename(args.img_path)
    image = cv2.imread(args.img_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    original_h = image.shape[0]
    original_w = image.shape[1]
    if sys.platform == "darwin":
            plt.switch_backend("TkAgg")
    plt.figure(figsize=(original_w/100, original_h/100))
    # Add subplot with no margin.
    plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
    plt.margins(0, 0)
    plt.gca().xaxis.set_major_locator(plt.NullLocator())
    plt.gca().yaxis.set_major_locator(plt.NullLocator())
    plt.imshow(image)
    if args.better_quality == True:
        if isinstance(annotations[0], torch.Tensor):
            annotations = np.array(annotations.cpu())
        for i, mask in enumerate(annotations):
            mask = cv2.morphologyEx(
                mask.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8)
            )
            annotations[i] = cv2.morphologyEx(
                mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8)
            )
    if args.device == "cpu":
        annotations = np.array(annotations)
        fast_show_mask(
            annotations,
            plt.gca(),
            random_color=mask_random_color,
            bbox=bbox,
            points=points,
            point_label=args.point_label,
            retinamask=args.retina,
            target_height=original_h,
            target_width=original_w,
        )
    else:
        if isinstance(annotations[0], np.ndarray):
            annotations = torch.from_numpy(annotations)
        fast_show_mask_gpu(
            annotations,
            plt.gca(),
            random_color=args.randomcolor,
            bbox=bbox,
            points=points,
            point_label=args.point_label,
            retinamask=args.retina,
            target_height=original_h,
            target_width=original_w,
        )
    if isinstance(annotations, torch.Tensor):
        annotations = annotations.cpu().numpy()
    if args.withContours == True:
        contour_all = []
        temp = np.zeros((original_h, original_w, 1))
        for i, mask in enumerate(annotations):
            if type(mask) == dict:
                mask = mask["segmentation"]
            annotation = mask.astype(np.uint8)
            if args.retina == False:
                annotation = cv2.resize(
                    annotation,
                    (original_w, original_h),
                    interpolation=cv2.INTER_NEAREST,
                )
            contours, hierarchy = cv2.findContours(
                annotation, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
            )
            for contour in contours:
                contour_all.append(contour)
        cv2.drawContours(temp, contour_all, -1, (255, 255, 255), 2)
        color = np.array([0 / 255, 0 / 255, 255 / 255, 0.8])
        contour_mask = temp / 255 * color.reshape(1, 1, -1)
        plt.imshow(contour_mask)

    save_path = args.output
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    plt.axis("off")
    fig = plt.gcf()
    plt.draw()
    
    try:
        buf = fig.canvas.tostring_rgb()
    except AttributeError:
        fig.canvas.draw()
        buf = fig.canvas.tostring_rgb()
    
    cols, rows = fig.canvas.get_width_height()
    img_array = np.fromstring(buf, dtype=np.uint8).reshape(rows, cols, 3)
    cv2.imwrite(os.path.join(save_path, result_name), cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR))


# CPU post process
def fast_show_mask(
    annotation,
    ax,
    random_color=False,
    bbox=None,
    points=None,
    point_label=None,
    retinamask=True,
    target_height=960,
    target_width=960,
):
    msak_sum = annotation.shape[0]
    height = annotation.shape[1]
    weight = annotation.shape[2]
    # 将annotation 按照面积 排序
    areas = np.sum(annotation, axis=(1, 2))
    sorted_indices = np.argsort(areas)
    annotation = annotation[sorted_indices]

    index = (annotation != 0).argmax(axis=0)
    if random_color == True:
        color = np.random.random((msak_sum, 1, 1, 3))
    else:
        color = np.ones((msak_sum, 1, 1, 3)) * np.array(
            [30 / 255, 144 / 255, 255 / 255]
        )
    transparency = np.ones((msak_sum, 1, 1, 1)) * 0.6
    visual = np.concatenate([color, transparency], axis=-1)
    mask_image = np.expand_dims(annotation, -1) * visual

    show = np.zeros((height, weight, 4))
    h_indices, w_indices = np.meshgrid(
        np.arange(height), np.arange(weight), indexing="ij"
    )
    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
    # 使用向量化索引更新show的值
    show[h_indices, w_indices, :] = mask_image[indices]
    if bbox is not None:
        x1, y1, x2, y2 = bbox
        ax.add_patch(
            plt.Rectangle(
                (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
            )
        )
    # draw point
    if points is not None:
        plt.scatter(
            [point[0] for i, point in enumerate(points) if point_label[i] == 1],
            [point[1] for i, point in enumerate(points) if point_label[i] == 1],
            s=20,
            c="y",
        )
        plt.scatter(
            [point[0] for i, point in enumerate(points) if point_label[i] == 0],
            [point[1] for i, point in enumerate(points) if point_label[i] == 0],
            s=20,
            c="m",
        )

    if retinamask == False:
        show = cv2.resize(
            show, (target_width, target_height), interpolation=cv2.INTER_NEAREST
        )
    ax.imshow(show)


def fast_show_mask_gpu(
    annotation,
    ax,
    random_color=False,
    bbox=None,
    points=None,
    point_label=None,
    retinamask=True,
    target_height=960,
    target_width=960,
):
    msak_sum = annotation.shape[0]
    height = annotation.shape[1]
    weight = annotation.shape[2]
    areas = torch.sum(annotation, dim=(1, 2))
    sorted_indices = torch.argsort(areas, descending=False)
    annotation = annotation[sorted_indices]
    # 找每个位置第一个非零值下标
    index = (annotation != 0).to(torch.long).argmax(dim=0)
    if random_color == True:
        color = torch.rand((msak_sum, 1, 1, 3)).to(annotation.device)
    else:
        color = torch.ones((msak_sum, 1, 1, 3)).to(annotation.device) * torch.tensor(
            [30 / 255, 144 / 255, 255 / 255]
        ).to(annotation.device)
    transparency = torch.ones((msak_sum, 1, 1, 1)).to(annotation.device) * 0.6
    visual = torch.cat([color, transparency], dim=-1)
    mask_image = torch.unsqueeze(annotation, -1) * visual
    # 按index取数，index指每个位置选哪个batch的数，把mask_image转成一个batch的形式
    show = torch.zeros((height, weight, 4)).to(annotation.device)
    h_indices, w_indices = torch.meshgrid(
        torch.arange(height), torch.arange(weight), indexing="ij"
    )
    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
    # 使用向量化索引更新show的值
    show[h_indices, w_indices, :] = mask_image[indices]
    show_cpu = show.cpu().numpy()
    if bbox is not None:
        x1, y1, x2, y2 = bbox
        ax.add_patch(
            plt.Rectangle(
                (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
            )
        )
    # draw point
    if points is not None:
        plt.scatter(
            [point[0] for i, point in enumerate(points) if point_label[i] == 1],
            [point[1] for i, point in enumerate(points) if point_label[i] == 1],
            s=20,
            c="y",
        )
        plt.scatter(
            [point[0] for i, point in enumerate(points) if point_label[i] == 0],
            [point[1] for i, point in enumerate(points) if point_label[i] == 0],
            s=20,
            c="m",
        )
    if retinamask == False:
        show_cpu = cv2.resize(
            show_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST
        )
    ax.imshow(show_cpu)


# clip
@torch.no_grad()
def retriev(
    model, preprocess, elements: [Image.Image], search_text: str, device
):
    preprocessed_images = [preprocess(image).to(device) for image in elements]
    tokenized_text = clip.tokenize([search_text]).to(device)
    stacked_images = torch.stack(preprocessed_images)
    image_features = model.encode_image(stacked_images)
    text_features = model.encode_text(tokenized_text)
    image_features /= image_features.norm(dim=-1, keepdim=True)
    text_features /= text_features.norm(dim=-1, keepdim=True)
    probs = 100.0 * image_features @ text_features.T
    return probs[:, 0].softmax(dim=0)


def crop_image(annotations, image_like):
    if isinstance(image_like, str):
        image = Image.open(image_like)
    else:
        image = image_like
    ori_w, ori_h = image.size
    mask_h, mask_w = annotations[0]["segmentation"].shape
    if ori_w != mask_w or ori_h != mask_h:
        image = image.resize((mask_w, mask_h))
    cropped_boxes = []
    cropped_images = []
    not_crop = []
    origin_id = []
    for _, mask in enumerate(annotations):
        if np.sum(mask["segmentation"]) <= 100:
            continue
        origin_id.append(_)
        bbox = get_bbox_from_mask(mask["segmentation"])  # mask 的 bbox
        cropped_boxes.append(segment_image(image, bbox))  # 保存裁剪的图片
        # cropped_boxes.append(segment_image(image,mask["segmentation"]))
        cropped_images.append(bbox)  # 保存裁剪的图片的bbox
    return cropped_boxes, cropped_images, not_crop, origin_id, annotations


def box_prompt(masks, bbox, target_height, target_width):
    h = masks.shape[1]
    w = masks.shape[2]
    if h != target_height or w != target_width:
        bbox = [
            int(bbox[0] * w / target_width),
            int(bbox[1] * h / target_height),
            int(bbox[2] * w / target_width),
            int(bbox[3] * h / target_height),
        ]
    bbox[0] = round(bbox[0]) if round(bbox[0]) > 0 else 0
    bbox[1] = round(bbox[1]) if round(bbox[1]) > 0 else 0
    bbox[2] = round(bbox[2]) if round(bbox[2]) < w else w
    bbox[3] = round(bbox[3]) if round(bbox[3]) < h else h

    # IoUs = torch.zeros(len(masks), dtype=torch.float32)
    bbox_area = (bbox[3] - bbox[1]) * (bbox[2] - bbox[0])

    masks_area = torch.sum(masks[:, bbox[1] : bbox[3], bbox[0] : bbox[2]], dim=(1, 2))
    orig_masks_area = torch.sum(masks, dim=(1, 2))

    union = bbox_area + orig_masks_area - masks_area
    IoUs = masks_area / union
    max_iou_index = torch.argmax(IoUs)

    return masks[max_iou_index].cpu().numpy(), max_iou_index


def point_prompt(masks, points, point_label, target_height, target_width):  # numpy 处理
    h = masks[0]["segmentation"].shape[0]
    w = masks[0]["segmentation"].shape[1]
    if h != target_height or w != target_width:
        points = [
            [int(point[0] * w / target_width), int(point[1] * h / target_height)]
            for point in points
        ]
    onemask = np.zeros((h, w))
    masks = sorted(masks, key=lambda x: x['area'], reverse=True)
    for i, annotation in enumerate(masks):
        if type(annotation) == dict:
            mask = annotation['segmentation']
        else:
            mask = annotation
        for i, point in enumerate(points):
            if mask[point[1], point[0]] == 1 and point_label[i] == 1:
                onemask[mask] = 1
            if mask[point[1], point[0]] == 1 and point_label[i] == 0:
                onemask[mask] = 0
    onemask = onemask >= 1
    return onemask, 0


def text_prompt(annotations, text, img_path, device, wider=False, threshold=0.9):
    cropped_boxes, cropped_images, not_crop, origin_id, annotations_ = crop_image(
        annotations, img_path
    )
    clip_model, preprocess = clip.load("ViT-B/32", device=device)
    scores = retriev(
        clip_model, preprocess, cropped_boxes, text, device=device
    )
    max_idx = scores.argsort()
    max_idx = max_idx[-1]
    max_idx = origin_id[int(max_idx)]

    # find the biggest mask which contains the mask with max score
    if wider:
        mask0 = annotations_[max_idx]["segmentation"]
        area0 = np.sum(mask0)
        areas = [(i, np.sum(mask["segmentation"])) for i, mask in enumerate(annotations_) if i in origin_id]
        areas = sorted(areas, key=lambda area: area[1], reverse=True)
        indices = [area[0] for area in areas]
        for index in indices:
            if index == max_idx or np.sum(annotations_[index]["segmentation"] & mask0) / area0 > threshold:
                max_idx = index
                break

    return annotations_[max_idx]["segmentation"], max_idx

DocImgTranslation/FastSAM/utils/tools.py
tianclll/Ace-Translate
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import cv2
import torch


def fast_process(
    annotations,
    image,
    device,
    scale,
    better_quality=False,
    mask_random_color=True,
    bbox=None,
    use_retina=True,
    withContours=True,
):
    if isinstance(annotations[0], dict):
        annotations = [annotation['segmentation'] for annotation in annotations]

    original_h = image.height
    original_w = image.width
    if better_quality:
        if isinstance(annotations[0], torch.Tensor):
            annotations = np.array(annotations.cpu())
        for i, mask in enumerate(annotations):
            mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8))
            annotations[i] = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8))
    if device == 'cpu':
        annotations = np.array(annotations)
        inner_mask = fast_show_mask(
            annotations,
            plt.gca(),
            random_color=mask_random_color,
            bbox=bbox,
            retinamask=use_retina,
            target_height=original_h,
            target_width=original_w,
        )
    else:
        if isinstance(annotations[0], np.ndarray):
            annotations = torch.from_numpy(annotations)
        inner_mask = fast_show_mask_gpu(
            annotations,
            plt.gca(),
            random_color=mask_random_color,
            bbox=bbox,
            retinamask=use_retina,
            target_height=original_h,
            target_width=original_w,
        )
    if isinstance(annotations, torch.Tensor):
        annotations = annotations.cpu().numpy()

    if withContours:
        contour_all = []
        temp = np.zeros((original_h, original_w, 1))
        for i, mask in enumerate(annotations):
            if type(mask) == dict:
                mask = mask['segmentation']
            annotation = mask.astype(np.uint8)
            if use_retina == False:
                annotation = cv2.resize(
                    annotation,
                    (original_w, original_h),
                    interpolation=cv2.INTER_NEAREST,
                )
            contours, _ = cv2.findContours(annotation, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
            for contour in contours:
                contour_all.append(contour)
        cv2.drawContours(temp, contour_all, -1, (255, 255, 255), 2 // scale)
        color = np.array([0 / 255, 0 / 255, 255 / 255, 0.9])
        contour_mask = temp / 255 * color.reshape(1, 1, -1)

    image = image.convert('RGBA')
    overlay_inner = Image.fromarray((inner_mask * 255).astype(np.uint8), 'RGBA')
    image.paste(overlay_inner, (0, 0), overlay_inner)

    if withContours:
        overlay_contour = Image.fromarray((contour_mask * 255).astype(np.uint8), 'RGBA')
        image.paste(overlay_contour, (0, 0), overlay_contour)

    return image


# CPU post process
def fast_show_mask(
    annotation,
    ax,
    random_color=False,
    bbox=None,
    retinamask=True,
    target_height=960,
    target_width=960,
):
    mask_sum = annotation.shape[0]
    height = annotation.shape[1]
    weight = annotation.shape[2]
    # 将annotation 按照面积 排序
    areas = np.sum(annotation, axis=(1, 2))
    sorted_indices = np.argsort(areas)[::1]
    annotation = annotation[sorted_indices]

    index = (annotation != 0).argmax(axis=0)
    if random_color:
        color = np.random.random((mask_sum, 1, 1, 3))
    else:
        color = np.ones((mask_sum, 1, 1, 3)) * np.array([30 / 255, 144 / 255, 255 / 255])
    transparency = np.ones((mask_sum, 1, 1, 1)) * 0.6
    visual = np.concatenate([color, transparency], axis=-1)
    mask_image = np.expand_dims(annotation, -1) * visual

    mask = np.zeros((height, weight, 4))

    h_indices, w_indices = np.meshgrid(np.arange(height), np.arange(weight), indexing='ij')
    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))

    mask[h_indices, w_indices, :] = mask_image[indices]
    if bbox is not None:
        x1, y1, x2, y2 = bbox
        ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor='b', linewidth=1))

    if not retinamask:
        mask = cv2.resize(mask, (target_width, target_height), interpolation=cv2.INTER_NEAREST)

    return mask


def fast_show_mask_gpu(
    annotation,
    ax,
    random_color=False,
    bbox=None,
    retinamask=True,
    target_height=960,
    target_width=960,
):
    device = annotation.device
    mask_sum = annotation.shape[0]
    height = annotation.shape[1]
    weight = annotation.shape[2]
    areas = torch.sum(annotation, dim=(1, 2))
    sorted_indices = torch.argsort(areas, descending=False)
    annotation = annotation[sorted_indices]
    # 找每个位置第一个非零值下标
    index = (annotation != 0).to(torch.long).argmax(dim=0)
    if random_color:
        color = torch.rand((mask_sum, 1, 1, 3)).to(device)
    else:
        color = torch.ones((mask_sum, 1, 1, 3)).to(device) * torch.tensor(
            [30 / 255, 144 / 255, 255 / 255]
        ).to(device)
    transparency = torch.ones((mask_sum, 1, 1, 1)).to(device) * 0.6
    visual = torch.cat([color, transparency], dim=-1)
    mask_image = torch.unsqueeze(annotation, -1) * visual
    # 按index取数，index指每个位置选哪个batch的数，把mask_image转成一个batch的形式
    mask = torch.zeros((height, weight, 4)).to(device)
    h_indices, w_indices = torch.meshgrid(torch.arange(height), torch.arange(weight))
    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
    # 使用向量化索引更新show的值
    mask[h_indices, w_indices, :] = mask_image[indices]
    mask_cpu = mask.cpu().numpy()
    if bbox is not None:
        x1, y1, x2, y2 = bbox
        ax.add_patch(
            plt.Rectangle(
                (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
            )
        )
    if not retinamask:
        mask_cpu = cv2.resize(
            mask_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST
        )
    return mask_cpu

DocImgTranslation/FastSAM/utils/tools_gradio.py
tianclll/Ace-Translate
import cv2
import numpy as np
from FastSAM import Inference


def order_points(pts):
    # 一共4个坐标点
    rect = np.zeros((4, 2), dtype="float32")
    # 按顺序找到对应坐标0123分别是 左上，右上，右下，左下
    # 计算左上，右下
    s = pts.sum(axis = 1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]
    # 计算右上和左下
    diff = np.diff(pts, axis = 1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]
    return rect


def four_point_transform(image, pts):
    # 获取输入坐标点
    rect = order_points(pts)
    (tl, tr, br, bl) = rect

    # 计算输入的w和h值
    widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
    widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
    maxWidth = max(int(widthA), int(widthB))

    heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
    heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
    maxHeight = max(int(heightA), int(heightB))

    # 变换后对应坐标位置
    dst = np.array([
        [0, 0],
        [maxWidth - 1, 0],
        [maxWidth - 1, maxHeight - 1],
        [0, maxHeight - 1]], dtype="float32")

    # 计算变换矩阵
    M = cv2.getPerspectiveTransform(rect, dst)
    warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))
    # 返回变换后结果
    return warped


def image_segmentation(img_path):
    img = cv2.imread(img_path)
    ann = Inference.main(img_path)
    ann = ann.numpy()
    ann = ann.astype('uint8')
    # ann = np.reshape(ann, (ann.shape[1],ann.shape[2]))
    ann[ann == 1] = 255
    contours = []
    # print(ann.shape)
    for i in ann:
        # print(i.shape)
        kernel = np.ones((20, 20), np.uint8)
        i = cv2.erode(i, kernel, iterations=1)
        contour, hierarchy = cv2.findContours(i, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        contours.append(contour[0])
    docCnt = None
    if len(contours) > 0:
        # 根据轮廓大小进行排序
        cnts = sorted(contours, key=cv2.contourArea, reverse=True)
        # 遍历每一个轮廓
        for c in cnts:
            # 近似
            peri = cv2.arcLength(c, True)
            approx = cv2.approxPolyDP(c, 0.02 * peri, True)
            # 准备做透视变换
            if len(approx) == 4:
                if approx[0][0][0] == 0 and approx[0][0][1] == 0:
                    continue
                docCnt = approx
                break

    warped = four_point_transform(img, docCnt.reshape(4, 2))
    return warped
if __name__ == '__main__':
    image_segmentation("./16.jpg")

DocImgTranslation/SegmentImage.py
tianclll/Ace-Translate
import cv2
import numpy as np


# Gamma校正 fGamaa=0.45是常用值
def gamma_correction(src, fGamma):
    # 构建查找表
    lut = np.array([np.uint8(pow(i / 255.0, fGamma) * 255) for i in range(256)])
    dst = cv2.LUT(src, lut)
    return dst


def image_enhancement(image_data):
    # image_data = cv2.cvtColor(image_data, cv2.COLOR_BGR2GRAY)
    # 划分算法
    # 如果混合色与基色相同则结果色为白色
    # 如混合色为白色则结果色为基色不变
    # 如混合色为黑色则结果色为白色
    src = image_data.astype(np.float32) / 255.0
    gauss = cv2.GaussianBlur(src, (101, 101), 0)
    dst = src / gauss
    dst = np.clip(dst * 255, 0, 255).astype(np.uint8)

    # gamma变换
    matGamma = gamma_correction(dst.copy(), 1.5)

    # 显示最终结果
    cv2.imwrite("DocImgTranslation/imageEnhancement.jpg", matGamma)
    return matGamma

if __name__ == "__main__":
    P = cv2.imread('12.jpg', cv2.IMREAD_COLOR)
    gray_P = cv2.cvtColor(P, cv2.COLOR_BGR2GRAY)
    P_shape = gray_P.shape

    image_enhancement(gray_P)


DocImgTranslation/imageEnhancement.py
tianclll/Ace-Translate
#!/usr/bin/env python
import requests
import json
import base64
import numpy as np
import cv2
def get_file_content(filePath):
    with open(filePath, 'rb') as fp:
        return fp.read()

class CommonOcr(object):
    def __init__(self, img_path):
        # 请登录后前往 “工作台-账号设置-开发者信息” 查看 x-ti-app-id
        # 示例代码中 x-ti-app-id 非真实数据
        self._app_id = 'a1ad6dc8a0a9f5da5487cb5a1a6e5bcf'
        # 请登录后前往 “工作台-账号设置-开发者信息” 查看 x-ti-secret-code
        # 示例代码中 x-ti-secret-code 非真实数据
        self._secret_code = '9ec1a530b4c985ed431891ca83e19acd'
        self._img_path = img_path

    def recognize(self):
        # 文档图像切边矫正
        url = 'https://api.textin.com/ai/service/v1/dewarp'
        head = {}
        try:
            image = get_file_content(self._img_path)
            head['x-ti-app-id'] = self._app_id
            head['x-ti-secret-code'] = self._secret_code
            result = requests.post(url, data=image, headers=head)
            return result.text
        except Exception as e:
            return e
def image_segmentation(image_path):
    response = CommonOcr(image_path)
    result = response.recognize()
    result = json.loads(result)["result"]["image"]
    img_binary = base64.b64decode(result)
    # Convert binary to numpy array
    img_np = np.fromstring(img_binary, np.uint8)
    # Convert numpy array to image
    img = cv2.imdecode(img_np, cv2.IMREAD_COLOR)
    return img
if __name__ == "__main__":
    response = CommonOcr(r'/Users/liuhongdi/计算机设计大赛/OCR图片文字识别/数据/JPG/000022.JPG')
    result = response.recognize()
    result = json.loads(result)["result"]["image"]

    img_binary = base64.b64decode(result)

    # Convert binary to numpy array
    img_np = np.fromstring(img_binary, np.uint8)

    # Convert numpy array to image
    img = cv2.imdecode(img_np, cv2.IMREAD_COLOR)
DocImgTranslation/imageSegmentation1.py
tianclll/Ace-Translate
import os
import argparse
from structure import Structure
from imageEnhancement import image_enhancement
from SegmentImage import image_segmentation
from paddleocr.ppocr.utils.logging import get_logger

logger = get_logger()
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", help="path to the input image")
ap.add_argument("-d","--image_dir")
ap.add_argument("-o","--output")
args = vars(ap.parse_args())

if args['image'] != None and args['image_dir'] == None:
    logger.info('开始分割图像...')
    img = image_segmentation(args['image'])
    #img = cv2.imread(args['image'],0)
    logger.info('分割图像已完成!')
    logger.info('开始图像增强...')
    img = image_enhancement(img)
    logger.info('图像增强已完成!')
    logger.info('开始版面分析并转为Word...')
    Structure(img,args['image'],args['output'],"ch2en")
    logger.info('已处理完成!')
elif args['image'] == None and args['image_dir'] != None:
    image_dir = args['image_dir']
    if os.path.exists(image_dir):
        contents = os.listdir(image_dir)
        for i, image in enumerate(contents):
            if image[-3:] not in ['jpg', 'JPG', 'png', 'PNG', 'peg']:
                contents.pop(i)
        logger.info('检测到{}张图片!'.format(len(contents)))
        for i,image in enumerate(contents):
            logger.info('正在处理第{}张图片...'.format(i+1))
            # image = cv2.imread(os.path.join(image_dir,i))
            logger.info('开始分割图像...')
            img = image_segmentation(os.path.join(image_dir,image))
            logger.info('分割图像已完成!')
            logger.info('开始图像增强...')
            img = image_enhancement(img)
            logger.info('图像增强已完成!')
            logger.info('开始版面分析并转为Word...')
            Structure(img, os.path.join(image_dir,image),args['output'])
            logger.info('第{}张图片已完成!'.format(i+1))
            print()
            if i == len(contents)-1:
                logger.info('已全部处理完成!')
                break
    else:
        print('请输入正确的路径!!!')


DocImgTranslation/ocr_image_ch2en.py
tianclll/Ace-Translate
import os
import argparse
from structure import Structure
from imageEnhancement import image_enhancement
from SegmentImage import image_segmentation
from paddleocr.ppocr.utils.logging import get_logger

logger = get_logger()
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", help="path to the input image")
ap.add_argument("-d","--image_dir")
ap.add_argument("-o","--output")
args = vars(ap.parse_args())

if args['image'] != None and args['image_dir'] == None:
    logger.info('开始分割图像...')
    img = image_segmentation(args['image'])
    #img = cv2.imread(args['image'],0)
    logger.info('分割图像已完成!')
    logger.info('开始图像增强...')
    img = image_enhancement(img)
    logger.info('图像增强已完成!')
    logger.info('开始版面分析并转为Word...')
    Structure(img,args['image'],args['output'],"en2ch")
    logger.info('已处理完成!')
elif args['image'] == None and args['image_dir'] != None:
    image_dir = args['image_dir']
    if os.path.exists(image_dir):
        contents = os.listdir(image_dir)
        for i, image in enumerate(contents):
            if image[-3:] not in ['jpg', 'JPG', 'png', 'PNG', 'peg']:
                contents.pop(i)
        logger.info('检测到{}张图片!'.format(len(contents)))
        for i,image in enumerate(contents):
            logger.info('正在处理第{}张图片...'.format(i+1))
            # image = cv2.imread(os.path.join(image_dir,i))
            logger.info('开始分割图像...')
            img = image_segmentation(os.path.join(image_dir,image))
            logger.info('分割图像已完成!')
            logger.info('开始图像增强...')
            img = image_enhancement(img)
            logger.info('图像增强已完成!')
            logger.info('开始版面分析并转为Word...')
            Structure(img, os.path.join(image_dir,image),args['output'])
            logger.info('第{}张图片已完成!'.format(i+1))
            print()
            if i == len(contents)-1:
                logger.info('已全部处理完成!')
                break
    else:
        print('请输入正确的路径!!!')


DocImgTranslation/ocr_image_en2ch.py
tianclll/Ace-Translate
import os
import cv2
import sys
sys.path.append(os.getcwd())
import my_utils
import numpy as np
from bs4 import BeautifulSoup
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from PIL import Image,ImageFont,ImageDraw
from paddleocr import PPStructure,save_structure_res,PaddleOCR
from paddleocr.ppstructure.recovery.recovery_to_doc import sorted_layout_boxes, convert_info_docx

def image_translate(img,text_region,text):
    x_coords = [text_region[0][0], text_region[1][0], text_region[2][0], text_region[3][0]]
    y_coords = [text_region[0][1],text_region[1][1],text_region[2][1],text_region[3][1]]
    x1, x2 = int(min(x_coords)), int(max(x_coords))
    y1, y2 = int(min(y_coords)), int(max(y_coords))
    img[y1:y2, x1:x2] = (255, 255, 255)  # 将区域填充为背景色
    bbox_width = max(x_coords) - min(x_coords)
    bbox_height = max(y_coords) - min(y_coords)
    font_size = 12
    # 将NumPy数组转换为PIL Image对象
    image_pil = Image.fromarray(img)
    draw = ImageDraw.Draw(image_pil)
    font = ImageFont.truetype("static/simfang.ttf", font_size)
    while draw.textsize(text, font=font)[0] < bbox_width and draw.textsize(text, font=font)[
        1] < bbox_height:
        font_size += 1
        font = ImageFont.truetype("static/simfang.ttf", font_size)
    text_x = text_region[0][0]
    text_y = text_region[0][1]
    font_color = (0, 0, 0)
    draw.text((text_x, text_y), text, font=font, fill=font_color)
    # image_pil.save("output_image.jpg")
    # 将修改后的PIL Image对象转回OpenCV格式
    img = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
    return img
def table_translate(html_table,translate_model):
    # 解析 HTML
    soup = BeautifulSoup(html_table, 'html.parser')
    # 提取并翻译表格的第一行（表头）
    original_header_row = soup.find('tr')
    translated_header_row = '<tr>'
    for cell in original_header_row.find_all('td'):
        original_text = cell.text.strip()
        translated_text = translate_model(original_text)[0]['translation_text']
        translated_text = my_utils.do_sentence(translated_text)
        translated_cell = f'<td>{translated_text}</td>'
        translated_header_row += translated_cell
    translated_header_row += '</tr>'

    # 提取并翻译表格的内容行
    translated_table = soup.new_tag('table')  # 创建一个新的表格用于存储翻译后的文本
    translated_table.append(BeautifulSoup(translated_header_row, 'html.parser'))  # 将翻译后的表头行插入到表格中
    for row in soup.find_all('tr')[1:]:
        translated_row = '<tr>'  # 创建新的行
        for cell in row.find_all('td'):
            original_text = cell.text.strip()  # 提取原始文本
            translated_text = translate_model(original_text)[0]['translation_text']
            translated_text = my_utils.do_sentence(translated_text)
            translated_cell = f'<td>{translated_text}</td>'  # 创建新的单元格
            translated_row += translated_cell  # 将单元格添加到行中
        translated_row += '</tr>'
        translated_table.append(BeautifulSoup(translated_row, 'html.parser'))  # 将行添加到表格中

    return translated_table
def Structure(img,img_path,save_path,mode):
    # 加载翻译模型
    if mode == "en2ch":
        model_path = 'models/translate/en-zh'
    # translate_model = hub.Module(name='transformer_zh-en', beam_size=5)
    # translate_model = hub.Module(name='baidu_translate')
    #     translate_model = hub.Module(name='tfbasemt_enzh', beam_size=5)
        # English image
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
        translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
        table_engine = PPStructure(show_log=False, recovery=True, image_orientation=True, lang='en',cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/en/en_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/en/en_PP-OCRv4_rec_infer',
                               layout_model_dir='models/ocr/layout/picodet_lcnet_x1_0_fgd_layout_infer',table_model_dir='models/ocr/table/en_ppstructure_mobile_v2.0_SLANet_infer')
        ocr = PaddleOCR(use_angle_cls=True, lang='ch',cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/en/en_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/en/en_PP-OCRv4_rec_infer')
    else:
        model_path = 'models/translate/zh-en'
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
        translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
        # Chinese image
        table_engine = PPStructure(show_log=False, recovery=True, image_orientation=True, layout_model_dir="models/ocr/layout/CDLA dict",
                                   structure_version="PP-StructureV2",cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/ch/ch_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/ch/ch_PP-OCRv4_rec_infer',
                                   table_model_dir='models/ocr/table/ch_ppstructure_mobile_v2.0_SLANet_infer')
        ocr = PaddleOCR(use_angle_cls=True, lang='ch',cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/ch/ch_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/ch/ch_PP-OCRv4_rec_infer')

    save_folder = os.path.join(save_path,'./output')
    result = table_engine(img)
    for i in range(len(result)):
        if result[i]["type"] != "figure" and result[i]["type"] != "table" and result[i]["type"] != "equation":
            # 本地中译英模型
            for j in result[i]["res"]:
                # print(j)
                content = j["text"]
                translate_text = translate_model(content)[0]['translation_text']
                translate_text = my_utils.do_sentence(translate_text)
                j["text"] = translate_text
        elif result[i]["type"] == "figure":
            roi_img = result[i]["img"]
            res = ocr.ocr(roi_img, cls=True)
            for j in res[0]:
                content = j[1][0]
                translate_text = translate_model(content)[0]['translation_text']
                translate_text = my_utils.do_sentence(translate_text)
                text_region = j[0]
                roi_img = image_translate(roi_img, text_region, translate_text)
            result[i]['img'] = roi_img
        elif result[i]["type"] == "table":
            result[i]["res"]["html"] = "{}".format(table_translate(result[i]["res"]["html"], translate_model))
    save_structure_res(result, save_folder, os.path.basename(img_path).split('.')[0])

    # for line in result:
    #     line.pop('img')
    #     # print(line)
    h, w, _ = img.shape
    res = sorted_layout_boxes(result, w)
    convert_info_docx(img, res, save_folder, os.path.basename(img_path).split('.')[0])
DocImgTranslation/structure.py
tianclll/Ace-Translate
import os
import sys
import pandas as pd
from bs4 import BeautifulSoup
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
def table_translate(html_table,translate_model):
    # 解析 HTML
    soup = BeautifulSoup(html_table, 'html.parser')
    # 提取并翻译表格的第一行（表头）
    original_header_row = soup.find('tr')
    translated_header_row = '<tr>'
    for cell in original_header_row.find_all('td'):
        original_text = cell.text.strip()
        translated_text = translate_model(original_text)[0]['translation_text']
        translated_cell = f'<td>{translated_text}</td>'
        translated_header_row += translated_cell
    translated_header_row += '</tr>'

    # 提取并翻译表格的内容行
    translated_table = soup.new_tag('table')  # 创建一个新的表格用于存储翻译后的文本
    translated_table.append(BeautifulSoup(translated_header_row, 'html.parser'))  # 将翻译后的表头行插入到表格中
    for row in soup.find_all('tr')[1:]:
        translated_row = '<tr>'  # 创建新的行
        for cell in row.find_all('td'):
            original_text = cell.text.strip()  # 提取原始文本
            translated_text = translate_model(original_text)[0]['translation_text']
            translated_cell = f'<td>{translated_text}</td>'  # 创建新的单元格
            translated_row += translated_cell  # 将单元格添加到行中
        translated_row += '</tr>'
        translated_table.append(BeautifulSoup(translated_row, 'html.parser'))  # 将行添加到表格中

    return translated_table
if __name__ == '__main__':
    savePath = sys.argv[2]
    # 1、excel地址
    excelPath = sys.argv[1]
    model_path = 'models/translate/zh-en'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
    # 读取 Excel 文件
    df = pd.read_excel(excelPath, sheet_name=0, header=None)
    html = df.to_html(index=False)
    html_string = table_translate(html,translate_model)
    # 创建 HTML 模板
    html_template = f"""
    <html>
      <head>
        <title>Table HTML Page</title>
      </head>
      <body>
        {html_string}
      </body>
    </html>
    """

    with open("FileTranslation/Excel/tmp.html", "w") as file:
        file.write(html_template)
    # 使用 pandas 读取表格
    df = pd.read_html("FileTranslation/Excel/tmp.html", encoding='utf-8')[0]
    print(df)
    # 将数据框保存为 Excel 文件
    excel_name = os.path.basename(excelPath)
    df.to_excel(os.path.join(savePath, excel_name), index=False, header=None)
FileTranslation/Excel/excel_translation_ch2en.py
tianclll/Ace-Translate
import os
import pandas as pd
import sys
sys.path.append(os.getcwd())
import my_utils
from bs4 import BeautifulSoup
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
def table_translate(html_table,translate_model):
    # 解析 HTML
    soup = BeautifulSoup(html_table, 'html.parser')
    # 提取并翻译表格的第一行（表头）
    original_header_row = soup.find('tr')
    translated_header_row = '<tr>'
    for cell in original_header_row.find_all('td'):
        original_text = cell.text.strip()
        translated_text = translate_model(original_text)[0]['translation_text']
        translated_text = my_utils.do_sentence(translated_text)
        translated_cell = f'<td>{translated_text}</td>'
        translated_header_row += translated_cell
    translated_header_row += '</tr>'

    # 提取并翻译表格的内容行
    translated_table = soup.new_tag('table')  # 创建一个新的表格用于存储翻译后的文本
    translated_table.append(BeautifulSoup(translated_header_row, 'html.parser'))  # 将翻译后的表头行插入到表格中
    for row in soup.find_all('tr')[1:]:
        translated_row = '<tr>'  # 创建新的行
        for cell in row.find_all('td'):
            original_text = cell.text.strip()  # 提取原始文本
            translated_text = translate_model(original_text)[0]['translation_text']
            translated_text = my_utils.do_sentence(translated_text)
            translated_cell = f'<td>{translated_text}</td>'  # 创建新的单元格
            translated_row += translated_cell  # 将单元格添加到行中
        translated_row += '</tr>'
        translated_table.append(BeautifulSoup(translated_row, 'html.parser'))  # 将行添加到表格中

    return translated_table
if __name__ == '__main__':
    savePath = sys.argv[2]
    # 1、excel地址
    excelPath = sys.argv[1]
    model_path = 'models/translate/en-zh'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
    # 读取 Excel 文件
    df = pd.read_excel(excelPath, sheet_name=0, header=None)
    html = df.to_html(index=False)
    html_string = table_translate(html,translate_model)
    # 创建 HTML 模板
    html_template = f"""
    <html>
      <head>
        <title>Table HTML Page</title>
      </head>
      <body>
        {html_string}
      </body>
    </html>
    """

    with open("FileTranslation/Excel/tmp.html", "w") as file:
        file.write(html_template)
    # 使用 pandas 读取表格
    df = pd.read_html("FileTranslation/Excel/tmp.html", encoding='utf-8')[0]
    print(df)
    # 将数据框保存为 Excel 文件
    excel_name = os.path.basename(excelPath)
    df.to_excel(os.path.join(savePath, excel_name), index=False, header=None)
FileTranslation/Excel/excel_translation_en2ch.py
tianclll/Ace-Translate
import os
import cv2
import numpy as np
from bs4 import BeautifulSoup
from PIL import Image,ImageFont,ImageDraw
from paddleocr import PPStructure,save_structure_res,PaddleOCR
from paddleocr.ppstructure.recovery.recovery_to_doc import sorted_layout_boxes, convert_info_docx
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from paddleocr.ppocr.utils.logging import get_logger
import sys
logger = get_logger()
def image_translate(img,text_region,text):
    x_coords = [text_region[0][0], text_region[1][0], text_region[2][0], text_region[3][0]]
    y_coords = [text_region[0][1],text_region[1][1],text_region[2][1],text_region[3][1]]
    x1, x2 = int(min(x_coords)), int(max(x_coords))
    y1, y2 = int(min(y_coords)), int(max(y_coords))
    img[y1:y2, x1:x2] = (255, 255, 255)  # 将区域填充为背景色
    bbox_width = max(x_coords) - min(x_coords)
    bbox_height = max(y_coords) - min(y_coords)
    font_size = 12
    # 将NumPy数组转换为PIL Image对象
    image_pil = Image.fromarray(img)
    draw = ImageDraw.Draw(image_pil)
    font = ImageFont.truetype("static/simfang.ttf", font_size)
    while draw.textsize(text, font=font)[0] < bbox_width and draw.textsize(text, font=font)[
        1] < bbox_height:
        font_size += 1
        font = ImageFont.truetype("static/simfang.ttf", font_size)
    text_x = text_region[0][0]
    text_y = text_region[0][1]
    font_color = (0, 0, 0)
    draw.text((text_x, text_y), text, font=font, fill=font_color)
    # image_pil.save("output_image.jpg")
    # 将修改后的PIL Image对象转回OpenCV格式
    img = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
    return img
def table_translate(html_table,translate_model):
    # 解析 HTML
    soup = BeautifulSoup(html_table, 'html.parser')
    # 提取并翻译表格的第一行（表头）
    original_header_row = soup.find('tr')
    translated_header_row = '<tr>'
    for cell in original_header_row.find_all('td'):
        original_text = cell.text.strip()
        translated_text = translate_model(original_text)[0]['translation_text']
        translated_cell = f'<td>{translated_text}</td>'
        translated_header_row += translated_cell
    translated_header_row += '</tr>'

    # 提取并翻译表格的内容行
    translated_table = soup.new_tag('table')  # 创建一个新的表格用于存储翻译后的文本
    translated_table.append(BeautifulSoup(translated_header_row, 'html.parser'))  # 将翻译后的表头行插入到表格中
    for row in soup.find_all('tr')[1:]:
        translated_row = '<tr>'  # 创建新的行
        for cell in row.find_all('td'):
            original_text = cell.text.strip()  # 提取原始文本
            translated_text = translate_model(original_text)[0]['translation_text']
            translated_cell = f'<td>{translated_text}</td>'  # 创建新的单元格
            translated_row += translated_cell  # 将单元格添加到行中
        translated_row += '</tr>'
        translated_table.append(BeautifulSoup(translated_row, 'html.parser'))  # 将行添加到表格中

    return translated_table
def img2doc(img,image_name,savePath):
    # 加载翻译模型
    #translate_model = hub.Module(name='transformer_zh-en', beam_size=5)
    # translate_model = hub.Module(name='baidu_translate')
    model_path = 'models/translate/zh-en'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
    #table_engine = PPStructure(show_log=False, recovery=True, image_orientation=True, layout_model_dir="CDLA dict",structure_version="PP-StructureV2")
    table_engine = PPStructure(show_log=False, recovery=True, image_orientation=True, layout_model_dir="models/ocr/layout/CDLA dict",
                                   structure_version="PP-StructureV2",cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/ch/ch_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/ch/ch_PP-OCRv4_rec_infer',
                                   table_model_dir='models/ocr/table/ch_ppstructure_mobile_v2.0_SLANet_infer')
    ocr = PaddleOCR(use_angle_cls=True, lang='ch',cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/ch/ch_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/ch/ch_PP-OCRv4_rec_infer')
    save_folder = os.path.join(savePath,'output/')
    result = table_engine(img)
    for i in range(len(result)):
        if result[i]["type"] != "figure" and result[i]["type"] != "table" and result[i]["type"] != "equation":
            # 本地中译英模型
            for j in result[i]["res"]:
                content = j["text"]
                translate_text = translate_model(content)[0]['translation_text']
                j["text"] = translate_text
            # 百度模型英译中
                # content = j["text"]
                # translate_text = translate_model.translate(content)
                # j["text"] = translate_text
        elif result[i]["type"] == "figure":
            roi_img = result[i]["img"]
            res = ocr.ocr(roi_img, cls=True)
            for j in res[0]:
                content = j[1][0]
                translate_text = translate_model(content)[0]['translation_text']
                text_region = j[0]
                roi_img = image_translate(roi_img,text_region, translate_text)
            result[i]['img'] = roi_img
        elif result[i]["type"] == "table":
            result[i]["res"]["html"] = "{}".format(table_translate(result[i]["res"]["html"],translate_model))
    save_structure_res(result, save_folder, image_name)
    h, w, _ = img.shape
    res = sorted_layout_boxes(result, w)
    convert_info_docx(img, res, save_folder, image_name)
    logger.info('已全部处理完成!')
if __name__ == '__main__':
    savePath = sys.argv[2]
    imgPath = sys.argv[1]
    img_name = os.path.basename(imgPath).split('.')[0]
    img = cv2.imread(imgPath)
    img2doc(img, img_name, savePath)
FileTranslation/Image/image_translation_ch2en.py
tianclll/Ace-Translate
import os
import cv2
import numpy as np
import sys
sys.path.append(os.getcwd())
import my_utils
from bs4 import BeautifulSoup
from PIL import Image,ImageFont,ImageDraw
from paddleocr import PPStructure,save_structure_res,PaddleOCR
from paddleocr.ppstructure.recovery.recovery_to_doc import sorted_layout_boxes, convert_info_docx
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from paddleocr.ppocr.utils.logging import get_logger
import sys
logger = get_logger()
def image_translate(img,text_region,text):
    x_coords = [text_region[0][0], text_region[1][0], text_region[2][0], text_region[3][0]]
    y_coords = [text_region[0][1],text_region[1][1],text_region[2][1],text_region[3][1]]
    x1, x2 = int(min(x_coords)), int(max(x_coords))
    y1, y2 = int(min(y_coords)), int(max(y_coords))
    img[y1:y2, x1:x2] = (255, 255, 255)  # 将区域填充为背景色
    bbox_width = max(x_coords) - min(x_coords)
    bbox_height = max(y_coords) - min(y_coords)
    font_size = 12
    # 将NumPy数组转换为PIL Image对象
    image_pil = Image.fromarray(img)
    draw = ImageDraw.Draw(image_pil)
    font = ImageFont.truetype("static/simfang.ttf", font_size)
    while draw.textsize(text, font=font)[0] < bbox_width and draw.textsize(text, font=font)[
        1] < bbox_height:
        font_size += 1
        font = ImageFont.truetype("static/simfang.ttf", font_size)
    text_x = text_region[0][0]
    text_y = text_region[0][1]
    font_color = (0, 0, 0)
    draw.text((text_x, text_y), text, font=font, fill=font_color)
    # image_pil.save("output_image.jpg")
    # 将修改后的PIL Image对象转回OpenCV格式
    img = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
    return img
def table_translate(html_table,translate_model):
    # 解析 HTML
    soup = BeautifulSoup(html_table, 'html.parser')
    # 提取并翻译表格的第一行（表头）
    original_header_row = soup.find('tr')
    translated_header_row = '<tr>'
    for cell in original_header_row.find_all('td'):
        original_text = cell.text.strip()
        translated_text = translate_model(original_text)[0]['translation_text']
        translated_text = my_utils.do_sentence(translated_text)
        translated_cell = f'<td>{translated_text}</td>'
        translated_header_row += translated_cell
    translated_header_row += '</tr>'

    # 提取并翻译表格的内容行
    translated_table = soup.new_tag('table')  # 创建一个新的表格用于存储翻译后的文本
    translated_table.append(BeautifulSoup(translated_header_row, 'html.parser'))  # 将翻译后的表头行插入到表格中
    for row in soup.find_all('tr')[1:]:
        translated_row = '<tr>'  # 创建新的行
        for cell in row.find_all('td'):
            original_text = cell.text.strip()  # 提取原始文本
            translated_text = translate_model(original_text)[0]['translation_text']
            translated_text = my_utils.do_sentence(translated_text)
            translated_cell = f'<td>{translated_text}</td>'  # 创建新的单元格
            translated_row += translated_cell  # 将单元格添加到行中
        translated_row += '</tr>'
        translated_table.append(BeautifulSoup(translated_row, 'html.parser'))  # 将行添加到表格中

    return translated_table
def img2doc(img,image_name,savePath):
    # 加载翻译模型
    #translate_model = hub.Module(name='transformer_zh-en', beam_size=5)
    # translate_model = hub.Module(name='baidu_translate')
    model_path = 'models/translate/en-zh'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
    #table_engine = PPStructure(show_log=False, recovery=True, image_orientation=True, layout_model_dir="CDLA dict",structure_version="PP-StructureV2")
    table_engine = PPStructure(show_log=False,recovery=True, lang='en',cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/en/en_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/en/en_PP-OCRv4_rec_infer',
                               layout_model_dir='models/ocr/layout/picodet_lcnet_x1_0_fgd_layout_infer',table_model_dir='models/ocr/table/en_ppstructure_mobile_v2.0_SLANet_infer')
    ocr = PaddleOCR(use_angle_cls=True, lang='en',cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/en/en_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/en/en_PP-OCRv4_rec_infer')
    save_folder = os.path.join(savePath,'output/')
    result = table_engine(img)
    for i in range(len(result)):
        if result[i]["type"] != "figure" and result[i]["type"] != "table" and result[i]["type"] != "equation":
            # 本地中译英模型
            for j in result[i]["res"]:
                content = j["text"]
                translate_text = translate_model(content)[0]['translation_text']
                translate_text = my_utils.do_sentence(translate_text)
                j["text"] = translate_text
            # 百度模型英译中
                # content = j["text"]
                # translate_text = translate_model.translate(content)
                # j["text"] = translate_text
        elif result[i]["type"] == "figure":
            roi_img = result[i]["img"]
            res = ocr.ocr(roi_img, cls=True)
            for j in res[0]:
                content = j[1][0]
                translate_text = translate_model(content)[0]['translation_text']
                translate_text = my_utils.do_sentence(translate_text)
                text_region = j[0]
                roi_img = image_translate(roi_img,text_region, translate_text)
            result[i]['img'] = roi_img
        elif result[i]["type"] == "table":
            result[i]["res"]["html"] = "{}".format(table_translate(result[i]["res"]["html"],translate_model))
    save_structure_res(result, save_folder, image_name)
    h, w, _ = img.shape
    res = sorted_layout_boxes(result, w)
    convert_info_docx(img, res, save_folder, image_name)
    logger.info('已全部处理完成!')
if __name__ == '__main__':
    savePath = sys.argv[2]
    imgPath = sys.argv[1]
    img_name = os.path.basename(imgPath).split('.')[0]
    img = cv2.imread(imgPath)
    img2doc(img, img_name, savePath)
FileTranslation/Image/image_translation_en2ch.py
tianclll/Ace-Translate
import os
import sys
import cv2
import numpy as np
import fitz  # fitz就是pip install PyMuPDF
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from bs4 import BeautifulSoup
from PIL import Image,ImageDraw,ImageFont
from paddleocr import PPStructure,save_structure_res,PaddleOCR
from paddleocr.ppstructure.recovery.recovery_to_doc import sorted_layout_boxes, convert_info_docx
from paddleocr.ppocr.utils.logging import get_logger

logger = get_logger()
def pyMuPDF_fitz(pdfPath):
    pdf_name = os.path.basename(pdfPath).split('.')[0]
    pdfDoc = fitz.open(pdfPath)
    imgs = []
    logger.info('检测PDF文档有共{}页'.format(pdfDoc.page_count))
    logger.info('正在处理...')
    for pg in range(0,pdfDoc.page_count):
        page = pdfDoc[pg]
        rotate = int(0)
        # 每个尺寸的缩放系数为1.3，这将为我们生成分辨率提高2.6的图像。
        # 此处若是不做设置，默认图片大小为：792X612, dpi=96
        zoom_x = 2  # (1.33333333-->1056x816)   (2-->1584x1224)
        zoom_y = 2
        mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
        pm = page.get_pixmap(matrix=mat, alpha=False)
        img = Image.frombytes("RGB", [pm.width, pm.height], pm.samples)
        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
        imgs.append(img)
        # if not os.path.exists(imagePath):  # 判断存放图片的文件夹是否存在
        #     os.makedirs(imagePath)  # 若图片文件夹不存在就创建
        # img.imwrite('./1.pjg',0)  # 将图片写入指定的文件夹内

    return imgs,pdf_name
def image_translate(img,text_region,text):
    x_coords = [text_region[0][0], text_region[1][0], text_region[2][0], text_region[3][0]]
    y_coords = [text_region[0][1],text_region[1][1],text_region[2][1],text_region[3][1]]
    x1, x2 = int(min(x_coords)), int(max(x_coords))
    y1, y2 = int(min(y_coords)), int(max(y_coords))
    img[y1:y2, x1:x2] = (255, 255, 255)  # 将区域填充为背景色
    bbox_width = max(x_coords) - min(x_coords)
    bbox_height = max(y_coords) - min(y_coords)
    font_size = 1
    # 将NumPy数组转换为PIL Image对象
    image_pil = Image.fromarray(img)
    draw = ImageDraw.Draw(image_pil)
    font = ImageFont.truetype("static/simfang.ttf", font_size)
    while draw.textsize(text, font=font)[0] < bbox_width and draw.textsize(text, font=font)[
        1] < bbox_height:
        font_size += 1
        font = ImageFont.truetype("static/simfang.ttf", font_size)
    text_x = text_region[0][0]
    text_y = text_region[0][1]
    font_color = (0, 0, 0)
    draw.text((text_x, text_y), text, font=font, fill=font_color)
    # image_pil.save("output_image.jpg")
    # 将修改后的PIL Image对象转回OpenCV格式
    img = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
    return img
def table_translate(html_table,translate_model):
    # 解析 HTML
    soup = BeautifulSoup(html_table, 'html.parser')
    # 提取并翻译表格的第一行（表头）
    original_header_row = soup.find('tr')
    translated_header_row = '<tr>'
    for cell in original_header_row.find_all('td'):
        original_text = cell.text.strip()
        translated_text = translate_model(original_text)[0]['translation_text']
        translated_cell = f'<td>{translated_text}</td>'
        translated_header_row += translated_cell
    translated_header_row += '</tr>'

    # 提取并翻译表格的内容行
    translated_table = soup.new_tag('table')  # 创建一个新的表格用于存储翻译后的文本
    translated_table.append(BeautifulSoup(translated_header_row, 'html.parser'))  # 将翻译后的表头行插入到表格中
    for row in soup.find_all('tr')[1:]:
        translated_row = '<tr>'  # 创建新的行
        for cell in row.find_all('td'):
            original_text = cell.text.strip()  # 提取原始文本
            translated_text = translate_model(original_text)[0]['translation_text']
            translated_cell = f'<td>{translated_text}</td>'  # 创建新的单元格
            translated_row += translated_cell  # 将单元格添加到行中
        translated_row += '</tr>'
        translated_table.append(BeautifulSoup(translated_row, 'html.parser'))  # 将行添加到表格中

    return translated_table
def img2doc(imgs,pdf_name,imagePath):
    # 加载翻译模型
    #translate_model = hub.Module(name='transformer_zh-en', beam_size=5)
    # translate_model = hub.Module(name='baidu_translate')
    model_path = 'models/translate/zh-en'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
    all_res = []
    table_engine = PPStructure(show_log=False, recovery=True, image_orientation=True, layout_model_dir="models/ocr/layout/CDLA dict",
                                   structure_version="PP-StructureV2",cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/ch/ch_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/ch/ch_PP-OCRv4_rec_infer',
                                   table_model_dir='models/ocr/table/ch_ppstructure_mobile_v2.0_SLANet_infer')
    # table_engine = PPStructure(show_log=False,recovery=True, lang='en')
    ocr = PaddleOCR(use_angle_cls=True, lang='ch',cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/ch/ch_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/ch/ch_PP-OCRv4_rec_infer')
    save_folder = os.path.join(imagePath,'output/')
    for index, img in enumerate(imgs):
        result = table_engine(img)
        for i in range(len(result)):
            if result[i]["type"] != "figure" and result[i]["type"] != "table" and result[i]["type"] != "equation":
                #本地中译英模型
                for j in result[i]["res"]:
                    content = j["text"]
                    translate_text = translate_model(content)[0]['translation_text']
                    j["text"] = translate_text
                #百度模型英译中
                    # content = j["text"]
                    # translate_text = translate_model.translate(content)
                    # j["text"] = translate_text
            elif result[i]["type"] == "figure":
                roi_img = result[i]["img"]
                res = ocr.ocr(roi_img, cls=True)
                for j in res[0]:
                    content = j[1][0]
                    translate_text = translate_model(content)[0]['translation_text']
                    text_region = j[0]
                    roi_img = image_translate(roi_img,text_region, translate_text)
                result[i]['img'] = roi_img
            elif result[i]["type"] == "table":
                result[i]["res"]["html"] = "{}".format(table_translate(result[i]["res"]["html"],translate_model))

        save_structure_res(result, save_folder, pdf_name)
        h, w, _ = img.shape
        res = sorted_layout_boxes(result, w)
        all_res += res
        logger.info('第{}页已处理完成...'.format(index+1))
    convert_info_docx(imgs, all_res, save_folder, pdf_name)
    logger.info('已全部处理完成!')
if __name__ == "__main__":
    imagePath = sys.argv[2]
    # 1、PDF地址
    pdfPath = sys.argv[1]
    # 2、需要储存图片的目录
    imgs,pdf_name= pyMuPDF_fitz(pdfPath)
    img2doc(imgs,pdf_name,imagePath)
FileTranslation/PDF/pdf_translation_ch2en.py
tianclll/Ace-Translate
import os
import cv2
import sys
sys.path.append(os.getcwd())
import my_utils
import numpy as np
import fitz  # fitz就是pip install PyMuPDF
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from bs4 import BeautifulSoup
from PIL import Image,ImageDraw,ImageFont
from paddleocr import PPStructure,save_structure_res,PaddleOCR
from paddleocr.ppstructure.recovery.recovery_to_doc import sorted_layout_boxes, convert_info_docx
from paddleocr.ppocr.utils.logging import get_logger

logger = get_logger()
def pyMuPDF_fitz(pdfPath):
    pdf_name = os.path.basename(pdfPath).split('.')[0]
    pdfDoc = fitz.open(pdfPath)
    imgs = []
    logger.info('检测PDF文档有共{}页'.format(pdfDoc.page_count))
    logger.info('正在处理...')
    for pg in range(0,pdfDoc.page_count):
        page = pdfDoc[pg]
        rotate = int(0)
        # 每个尺寸的缩放系数为1.3，这将为我们生成分辨率提高2.6的图像。
        # 此处若是不做设置，默认图片大小为：792X612, dpi=96
        zoom_x = 2  # (1.33333333-->1056x816)   (2-->1584x1224)
        zoom_y = 2
        mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
        pm = page.get_pixmap(matrix=mat, alpha=False)
        img = Image.frombytes("RGB", [pm.width, pm.height], pm.samples)
        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
        imgs.append(img)
        # if not os.path.exists(imagePath):  # 判断存放图片的文件夹是否存在
        #     os.makedirs(imagePath)  # 若图片文件夹不存在就创建
        # img.imwrite('./1.pjg',0)  # 将图片写入指定的文件夹内

    return imgs,pdf_name
def image_translate(img,text_region,text):
    x_coords = [text_region[0][0], text_region[1][0], text_region[2][0], text_region[3][0]]
    y_coords = [text_region[0][1],text_region[1][1],text_region[2][1],text_region[3][1]]
    x1, x2 = int(min(x_coords)), int(max(x_coords))
    y1, y2 = int(min(y_coords)), int(max(y_coords))
    img[y1:y2, x1:x2] = (255, 255, 255)  # 将区域填充为背景色
    bbox_width = max(x_coords) - min(x_coords)
    bbox_height = max(y_coords) - min(y_coords)
    font_size = 1
    # 将NumPy数组转换为PIL Image对象
    image_pil = Image.fromarray(img)
    draw = ImageDraw.Draw(image_pil)
    font = ImageFont.truetype("static/simfang.ttf", font_size)
    while draw.textsize(text, font=font)[0] < bbox_width and draw.textsize(text, font=font)[
        1] < bbox_height:
        font_size += 1
        font = ImageFont.truetype("static/simfang.ttf", font_size)
    text_x = text_region[0][0]
    text_y = text_region[0][1]
    font_color = (0, 0, 0)
    draw.text((text_x, text_y), text, font=font, fill=font_color)
    # image_pil.save("output_image.jpg")
    # 将修改后的PIL Image对象转回OpenCV格式
    img = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
    return img
def table_translate(html_table,translate_model):
    # 解析 HTML
    soup = BeautifulSoup(html_table, 'html.parser')
    # 提取并翻译表格的第一行（表头）
    original_header_row = soup.find('tr')
    translated_header_row = '<tr>'
    for cell in original_header_row.find_all('td'):
        original_text = cell.text.strip()
        translated_text = translate_model(original_text)[0]['translation_text']
        translated_text = my_utils.do_sentence(translated_text)
        translated_cell = f'<td>{translated_text}</td>'
        translated_header_row += translated_cell
    translated_header_row += '</tr>'

    # 提取并翻译表格的内容行
    translated_table = soup.new_tag('table')  # 创建一个新的表格用于存储翻译后的文本
    translated_table.append(BeautifulSoup(translated_header_row, 'html.parser'))  # 将翻译后的表头行插入到表格中
    for row in soup.find_all('tr')[1:]:
        translated_row = '<tr>'  # 创建新的行
        for cell in row.find_all('td'):
            original_text = cell.text.strip()  # 提取原始文本
            translated_text = translate_model(original_text)[0]['translation_text']
            translated_text = my_utils.do_sentence(translated_text)
            translated_cell = f'<td>{translated_text}</td>'  # 创建新的单元格
            translated_row += translated_cell  # 将单元格添加到行中
        translated_row += '</tr>'
        translated_table.append(BeautifulSoup(translated_row, 'html.parser'))  # 将行添加到表格中

    return translated_table
def img2doc(imgs,pdf_name,imagePath):
    # 加载翻译模型
    #translate_model = hub.Module(name='transformer_zh-en', beam_size=5)
    # translate_model = hub.Module(name='baidu_translate')
    model_path = 'models/translate/en-zh'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
    all_res = []
    # table_engine = PPStructure(show_log=False, recovery=True, layout_model_dir="models/CDLA dict", save_pdf=True, structure_version="PP-StructureV2")
    table_engine = PPStructure(show_log=False,recovery=True, lang='en',cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/en/en_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/en/en_PP-OCRv4_rec_infer',
                               layout_model_dir='models/ocr/layout/picodet_lcnet_x1_0_fgd_layout_infer',table_model_dir='models/ocr/table/en_ppstructure_mobile_v2.0_SLANet_infer')
    ocr = PaddleOCR(use_angle_cls=True, lang='en',cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/en/en_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/en/en_PP-OCRv4_rec_infer')
    save_folder = os.path.join(imagePath,'output/')
    for index, img in enumerate(imgs):
        result = table_engine(img)
        for i in range(len(result)):
            if result[i]["type"] != "figure" and result[i]["type"] != "table" and result[i]["type"] != "equation":
                #本地中译英模型
                for j in result[i]["res"]:
                    content = j["text"]
                    translate_text = translate_model(content)[0]['translation_text']
                    translate_text = my_utils.do_sentence(translate_text)
                    j["text"] = translate_text
                #百度模型英译中
                    # content = j["text"]
                    # translate_text = translate_model.translate(content)
                    # j["text"] = translate_text
            elif result[i]["type"] == "figure":
                roi_img = result[i]["img"]
                res = ocr.ocr(roi_img, cls=True)
                for j in res[0]:
                    content = j[1][0]
                    translate_text = translate_model(content)[0]['translation_text']
                    translate_text = my_utils.do_sentence(translate_text)
                    text_region = j[0]
                    roi_img = image_translate(roi_img,text_region, translate_text)
                result[i]['img'] = roi_img
            elif result[i]["type"] == "table":
                result[i]["res"]["html"] = "{}".format(table_translate(result[i]["res"]["html"],translate_model))

        save_structure_res(result, save_folder, pdf_name)
        h, w, _ = img.shape
        res = sorted_layout_boxes(result, w)
        all_res += res
        logger.info('第{}页已处理完成...'.format(index+1))
    convert_info_docx(imgs, all_res, save_folder, pdf_name)
    logger.info('已全部处理完成!')
if __name__ == "__main__":
    imagePath = sys.argv[2]
    # 1、PDF地址
    pdfPath = sys.argv[1]
    # 2、需要储存图片的目录
    imgs,pdf_name= pyMuPDF_fitz(pdfPath)
    img2doc(imgs,pdf_name,imagePath)
FileTranslation/PDF/pdf_translation_en2ch.py
tianclll/Ace-Translate
from pptx import Presentation
from time import sleep
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import os
import sys
if __name__ == '__main__':
    savePath = sys.argv[2]
    # 1、ppt地址
    pptPath = sys.argv[1]
    model_path = 'models/translate/zh-en'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
    prs = Presentation(pptPath)

    print("start")
    for ns, slide in enumerate(prs.slides):
        for nsh, shape in enumerate(slide.shapes):
            if not shape.has_text_frame:
                continue
            for np, paragraph in enumerate(shape.text_frame.paragraphs):
                for rs, run in enumerate(paragraph.runs):
                    str_in = run.text
                    str_out = translate_model(str_in)[0]['translation_text']
                    prs.slides[ns].shapes[nsh].text_frame.paragraphs[np].runs[rs].text = str_out
                    sleep(1.5)
                    print(np)
    ppt_name = os.path.basename(pptPath)
    print(ppt_name)
    prs.save(os.path.join(savePath,ppt_name))
    print(os.path.join(savePath,ppt_name))
FileTranslation/PPT/ppt_translation_ch2en.py
tianclll/Ace-Translate
from pptx import Presentation
from time import sleep
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import os
import sys
sys.path.append(os.getcwd())
import my_utils


if __name__ == '__main__':
    savePath = sys.argv[2]
    # 1、ppt地址
    pptPath = sys.argv[1]
    model_path = 'models/translate/en-zh'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
    prs = Presentation(pptPath)

    print("start")
    for ns, slide in enumerate(prs.slides):
        for nsh, shape in enumerate(slide.shapes):
            if not shape.has_text_frame:
                continue
            for np, paragraph in enumerate(shape.text_frame.paragraphs):
                for rs, run in enumerate(paragraph.runs):
                    str_in = run.text
                    str_out = translate_model(str_in)[0]['translation_text']
                    str_out = my_utils.do_sentence(str_out)
                    prs.slides[ns].shapes[nsh].text_frame.paragraphs[np].runs[rs].text = str_out
                    sleep(1.5)
    ppt_name = os.path.basename(pptPath)
    prs.save(os.path.join(savePath,ppt_name))
FileTranslation/PPT/ppt_translation_en2ch.py
tianclll/Ace-Translate
import sys
import os
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
if __name__ == '__main__':
    savePath = sys.argv[2]
    # 1、excel地址
    txtPath = sys.argv[1]
    model_path = 'models/translate/zh-en'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
    with open(txtPath,'r') as f:
        text = f.readlines()
        f.close()

    lines = []
    for line in text:
        line_clean = line.replace('\n','')
        translate_line = translate_model(line_clean)[0]['translation_text']
        translate_line = translate_line + '\n'
        lines.append(translate_line)

    txt_name = os.path.basename(txtPath)
    with open(os.path.join(savePath,txt_name),'w') as f:
        f.writelines(lines)
        f.close()
FileTranslation/TXT/txt_translation_ch2en.py
tianclll/Ace-Translate

import os
import sys
sys.path.append(os.getcwd())
import my_utils
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
if __name__ == '__main__':
    savePath = sys.argv[2]
    # 1、excel地址
    txtPath = sys.argv[1]
    model_path = 'models/translate/en-zh'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
    with open(txtPath,'r') as f:
        text = f.readlines()
        f.close()

    lines = []
    for line in text:
        line_clean = line.replace('\n','')
        translate_line = translate_model(line_clean)[0]['translation_text']
        translate_line = my_utils.do_sentence(translate_line)
        translate_line = translate_line + '\n'
        lines.append(translate_line)

    txt_name = os.path.basename(txtPath)
    with open(os.path.join(savePath,txt_name),'w') as f:
        f.writelines(lines)
        f.close()
FileTranslation/TXT/txt_translation_en2ch.py
tianclll/Ace-Translate
import sys
from docx import Document
import os
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
def get_paragraphs_text(path):
    """
    获取所有段落的文本
    :param path: word路径
    :return: list类型，如：
        ['Test', 'hello world', ...]
    """
    document = Document(path)
    all_paragraphs = document.paragraphs
    paragraphs_text = []
    for paragraph in all_paragraphs:
        # 拼接一个list,包括段落的结构和内容
        paragraphs_text.append([paragraph.style.name,paragraph.text])
    return paragraphs_text
if __name__ == '__main__':
    savePath = sys.argv[2]
    # 1、excel地址
    wordPath = sys.argv[1]
    # 提取文档信息
    text = get_paragraphs_text(wordPath)
    model_path = 'models/translate/zh-en'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
    new_document = Document()
    for item in text:
        src_style = item[0]
        src_texts = [item[1]]
        print('标题层级：', src_style)
        print(src_texts)
        n_best = 1  # 每个输入样本的输出候选句子数量
        trg_texts = [translate_model(src_texts[0])[0]['translation_text']]
        print(trg_texts)
        if trg_texts == ['N']:
            trg_texts = ['']
        if src_style == 'Title':
            new_document.add_heading(trg_texts, 0)
        elif src_style[:7:] == 'Heading':
            new_document.add_heading(trg_texts, level=int(src_style[-1]))
        else:
            new_document.add_paragraph(trg_texts)
    word_name = os.path.basename(wordPath)
    new_document.save(os.path.join(savePath,word_name))
FileTranslation/Word/word_translation_ch2en.py
tianclll/Ace-Translate
import sys
from docx import Document
import os
sys.path.append(os.getcwd())
import my_utils
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
def get_paragraphs_text(path):
    """
    获取所有段落的文本
    :param path: word路径
    :return: list类型，如：
        ['Test', 'hello world', ...]
    """
    document = Document(path)
    all_paragraphs = document.paragraphs
    paragraphs_text = []
    for paragraph in all_paragraphs:
        # 拼接一个list,包括段落的结构和内容
        paragraphs_text.append([paragraph.style.name,paragraph.text])
    return paragraphs_text
if __name__ == '__main__':
    savePath = sys.argv[2]
    # 1、excel地址
    wordPath = sys.argv[1]
    # 提取文档信息
    text = get_paragraphs_text(wordPath)
    model_path = 'models/translate/en-zh'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    translate_model = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
    new_document = Document()
    for item in text:
        src_style = item[0]
        src_texts = [item[1]]
        print('标题层级：', src_style)
        print(src_texts)
        n_best = 1  # 每个输入样本的输出候选句子数量
        trg_texts = [my_utils.do_sentence(translate_model(src_texts[0])[0]['translation_text'])]
        print(trg_texts)
        if trg_texts == ['N']:
            trg_texts = ['']
        if src_style == 'Title':
            new_document.add_heading(trg_texts, 0)
        elif src_style[:7:] == 'Heading':
            new_document.add_heading(trg_texts, level=int(src_style[-1]))
        else:
            new_document.add_paragraph(trg_texts)
    word_name = os.path.basename(wordPath)
    new_document.save(os.path.join(savePath,word_name))
FileTranslation/Word/word_translation_en2ch.py
tianclll/Ace-Translate
import configparser
import tkinter as tk

import pyperclip
import transformers
from pynput import keyboard
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

config = configparser.ConfigParser()
config.read("config.conf")
class TranslatorApp(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("全能翻译")
        self.geometry("300x100")

        self.init_ui()
        self.withdraw()  # 初始时隐藏窗口

    def init_ui(self):
        self.protocol("WM_DELETE_WINDOW", self.on_closing)  # 拦截窗口关闭事件

        # self.label = tk.Label(self, text="使用快捷键 Ctrl+Shift+F 进行截图并识别文字：")
        # self.label.pack(pady=10)

        self.frame = tk.Frame(self)
        self.frame.pack()

        self.text_frame = tk.Frame(self.frame)
        self.text_frame.pack(padx=10, pady=10)

        self.text_label = tk.Label(self.text_frame, text="翻译结果:")
        self.text_label.pack()

        self.text_box = tk.Text(self.text_frame, height=5, width=40)
        self.text_box.pack(fill=tk.BOTH, expand=True)
        translation_shortcut = config.get("shortcuts", "translation")
        self.listener = keyboard.GlobalHotKeys({
            translation_shortcut: self.text_ocr
        })
        self.listener.start()
    #
    # def get_selected_text(self):
    #     try:
    #         selected_text = subprocess.check_output(['pbpaste'], universal_newlines=True)
    #         return selected_text
    #     except subprocess.CalledProcessError:
    #         return ""
    def text_ocr(self):
        # keyboard1 = Controller()
        # 模拟按下 Ctrl + C 组合键
        # with keyboard1.pressed(Key.cmd):
        #     keyboard1.press('c')
        #     keyboard1.release('c')
        # time.sleep(0.5)
        # text = self.get_selected_text()
        text = pyperclip.paste()
        model_path = 'models/translate/zh-en'
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        translate_model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
        pipeline = transformers.pipeline("translation", model=translate_model, tokenizer=tokenizer)
        translate_text = pipeline(text)
        self.display_translation(translate_text[0]['translation_text'])
        print(translate_text[0]['translation_text'])
        self.deiconify()  # 在截图完成后显示窗口
    def display_translation(self, text):
        self.text_box.delete("1.0", tk.END)  # 清空文本框
        self.text_box.insert(tk.END, text)
    def on_closing(self):
        self.withdraw()  # 仅隐藏窗口，不退出程序
if __name__ == "__main__":
    app = TranslatorApp()
    app.mainloop()
HyperTranslation/hyper_translation_ch2en.py
tianclll/Ace-Translate
import configparser
import tkinter as tk
import sys
import os
sys.path.append(os.getcwd())
import my_utils
import pyperclip
import transformers
from pynput import keyboard
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
config = configparser.ConfigParser()
config.read("config.conf")
class TranslatorApp(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("全能翻译")
        self.geometry("300x100")

        self.init_ui()
        self.withdraw()  # 初始时隐藏窗口

    def init_ui(self):
        self.protocol("WM_DELETE_WINDOW", self.on_closing)  # 拦截窗口关闭事件

        # self.label = tk.Label(self, text="使用快捷键 Ctrl+Shift+F 进行截图并识别文字：")
        # self.label.pack(pady=10)

        self.frame = tk.Frame(self)
        self.frame.pack()

        self.text_frame = tk.Frame(self.frame)
        self.text_frame.pack(padx=10, pady=10)

        self.text_label = tk.Label(self.text_frame, text="翻译结果:")
        self.text_label.pack()

        self.text_box = tk.Text(self.text_frame, height=5, width=40)
        self.text_box.pack(fill=tk.BOTH, expand=True)
        translation_shortcut = config.get("shortcuts", "translation")
        self.listener = keyboard.GlobalHotKeys({
            translation_shortcut: self.text_ocr
        })
        self.listener.start()
    #
    # def get_selected_text(self):
    #     try:
    #         selected_text = subprocess.check_output(['pbpaste'], universal_newlines=True)
    #         return selected_text
    #     except subprocess.CalledProcessError:
    #         return ""
    def text_ocr(self):
        # keyboard1 = Controller()
        # 模拟按下 Ctrl + C 组合键
        # with keyboard1.pressed(Key.cmd):
        #     keyboard1.press('c')
        #     keyboard1.release('c')
        # time.sleep(0.5)
        # text = self.get_selected_text()
        text = pyperclip.paste()
        model_path = 'models/translate/en-zh'
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        translate_model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
        pipeline = transformers.pipeline("translation", model=translate_model, tokenizer=tokenizer)
        translate_text = pipeline(text)
        translate_text = my_utils.do_sentence(translate_text[0]['translation_text'])
        self.display_translation(translate_text)
        print(translate_text)
        self.deiconify()  # 在截图完成后显示窗口
    def display_translation(self, text):
        self.text_box.delete("1.0", tk.END)  # 清空文本框
        self.text_box.insert(tk.END, text)
    def on_closing(self):
        self.withdraw()  # 仅隐藏窗口，不退出程序
if __name__ == "__main__":
    app = TranslatorApp()
    app.mainloop()

HyperTranslation/hyper_translation_en2ch.py
tianclll/Ace-Translate
import os
import subprocess
import tkinter as tk
import configparser
from PIL import Image, ImageTk
from pynput import keyboard
from paddleocr import PaddleOCR
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import transformers
import platform

config = configparser.ConfigParser()
config.read("config.conf")


class ScreenShotTool(tk.Tk):
    def __init__(self):
        super().__init__()

        self.title("全能翻译")
        self.geometry("700x200")

        self.init_ui()
        self.withdraw()  # 初始时隐藏窗口

    def init_ui(self):
        self.protocol("WM_DELETE_WINDOW", self.on_closing)  # 拦截窗口关闭事件

        self.frame = tk.Frame(self)
        self.frame.pack()

        self.image_frame = tk.Frame(self.frame)
        self.image_frame.pack(side=tk.LEFT, padx=10, pady=10)

        self.text_frame = tk.Frame(self.frame)
        self.text_frame.pack(side=tk.RIGHT, padx=10, pady=10)
        self.image_label1 = tk.Label(self.image_frame, text="截图:")
        self.image_label1.pack()
        self.image_label = tk.Label(self.image_frame)
        self.image_label.pack()

        self.text_label = tk.Label(self.text_frame, text="翻译结果:")
        self.text_label.pack()

        self.text_box = tk.Text(self.text_frame, height=10, width=40)
        self.text_box.pack(fill=tk.BOTH, expand=True)
        screenshot_shortcut = config.get("shortcuts", "screenshot")
        self.listener = keyboard.GlobalHotKeys({
            screenshot_shortcut: self.capture_and_ocr
        })
        self.listener.start()

    def capture_and_ocr(self):
        ocr = PaddleOCR(use_angle_cls=True, lang='ch',cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/ch/ch_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/ch/ch_PP-OCRv4_rec_infer')
        model_path = 'models/translate/zh-en'
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
        pipeline = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
        screenshot_path = "temp_screenshot.png"

        system = platform.system()
        if system == "Windows":
            from PIL import ImageGrab
            subprocess.run(["snippingtool", "/clip"])
            clipboard_image = ImageGrab.grabclipboard()

            if clipboard_image is not None:
                # 保存图像到指定路径
                clipboard_image.save(screenshot_path)
        elif system == "Linux":
            subprocess.run(["gnome-screenshot", "-a", "-f", screenshot_path])
        elif system == "Darwin":
            subprocess.run(["screencapture", "-i", screenshot_path])
        else:
            print("Unsupported operating system.")
            return

        image = Image.open(screenshot_path)
        image.thumbnail((300, 300))
        photo = ImageTk.PhotoImage(image)
        self.image_label.config(image=photo)
        self.image_label.image = photo

        recognized_text = []
        results = ocr.ocr(screenshot_path, cls=True)
        for idx in range(len(results)):
            res = results[idx]
            for line in res:
                translate_text = pipeline(line[1][0])[0]['translation_text']
                recognized_text.append(translate_text)
        translate_text = "\n".join(recognized_text)
        print(translate_text)
        self.text_box.delete("1.0", tk.END)  # 清空文本框
        self.text_box.insert(tk.END, translate_text)

        os.remove(screenshot_path)  # 删除临时截图文件

        self.deiconify()  # 在截图完成后显示窗口

    def on_closing(self):
        self.withdraw()  # 仅隐藏窗口，不退出程序


if __name__ == "__main__":
    app = ScreenShotTool()
    app.mainloop()

ScreenshotTranslation/screenshot_translation_ch2en.py
tianclll/Ace-Translate
import os
import subprocess
import tkinter as tk
import configparser
from PIL import Image, ImageTk
from pynput import keyboard
from paddleocr import PaddleOCR
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import transformers
import platform
import sys
sys.path.append(os.getcwd())
import my_utils

config = configparser.ConfigParser()
config.read("config.conf")


class ScreenShotTool(tk.Tk):
    def __init__(self):
        super().__init__()

        self.title("全能翻译")
        self.geometry("700x200")

        self.init_ui()
        self.withdraw()  # 初始时隐藏窗口

    def init_ui(self):
        self.protocol("WM_DELETE_WINDOW", self.on_closing)  # 拦截窗口关闭事件

        self.frame = tk.Frame(self)
        self.frame.pack()

        self.image_frame = tk.Frame(self.frame)
        self.image_frame.pack(side=tk.LEFT, padx=10, pady=10)

        self.text_frame = tk.Frame(self.frame)
        self.text_frame.pack(side=tk.RIGHT, padx=10, pady=10)
        self.image_label1 = tk.Label(self.image_frame, text="截图:")
        self.image_label1.pack()
        self.image_label = tk.Label(self.image_frame)
        self.image_label.pack()

        self.text_label = tk.Label(self.text_frame, text="翻译结果:")
        self.text_label.pack()

        self.text_box = tk.Text(self.text_frame, height=10, width=40)
        self.text_box.pack(fill=tk.BOTH, expand=True)
        screenshot_shortcut = config.get("shortcuts", "screenshot")
        self.listener = keyboard.GlobalHotKeys({
            screenshot_shortcut: self.capture_and_ocr
        })
        self.listener.start()

    def capture_and_ocr(self):
        ocr = PaddleOCR(use_angle_cls=True, lang='en',cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/en/en_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/en/en_PP-OCRv4_rec_infer')
        model_path = 'models/translate/en-zh'
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
        pipeline = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
        screenshot_path = "temp_screenshot.png"

        system = platform.system()

        if system == "Windows":
            from PIL import ImageGrab
            subprocess.run(["snippingtool", "/clip"])
            clipboard_image = ImageGrab.grabclipboard()

            if clipboard_image is not None:
                # 保存图像到指定路径
                clipboard_image.save(screenshot_path)
        elif system == "Linux":
            subprocess.run(["gnome-screenshot", "-a", "-f", screenshot_path])
        elif system == "Darwin":
            subprocess.run(["screencapture", "-i", screenshot_path])
        else:
            print("Unsupported operating system.")
            return

        image = Image.open(screenshot_path)
        image.thumbnail((300, 300))
        photo = ImageTk.PhotoImage(image)
        self.image_label.config(image=photo)
        self.image_label.image = photo

        recognized_text = []
        results = ocr.ocr(screenshot_path, cls=True)
        for idx in range(len(results)):
            res = results[idx]
            for line in res:
                translate_text = pipeline(line[1][0])[0]['translation_text']
                translate_text = my_utils.do_sentence(translate_text)
                recognized_text.append(translate_text)
        translate_text = "\n".join(recognized_text)
        print(translate_text)
        self.text_box.delete("1.0", tk.END)  # 清空文本框
        self.text_box.insert(tk.END, translate_text)

        os.remove(screenshot_path)  # 删除临时截图文件

        self.deiconify()  # 在截图完成后显示窗口

    def on_closing(self):
        self.withdraw()  # 仅隐藏窗口，不退出程序


if __name__ == "__main__":
    app = ScreenShotTool()
    app.mainloop()

ScreenshotTranslation/screenshot_translation_en2ch.py
tianclll/Ace-Translate
#! /usr/bin/env python
import os
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from paddlespeech.cli.asr.infer import ASRExecutor
from paddlespeech.cli.text.infer import TextExecutor
from pydub import AudioSegment
from pydub.utils import make_chunks
import configparser
config = configparser.ConfigParser()
config.read("config.conf")
device = config.get("devices", "device")
# 输入音频文件的路径（例如 MP3、FLAC、OGG 等）

def translate(input_audio_file):
    # 指定输出的 WAV 文件路径
    output_wav_file = "SpeechTranslation/output_audio.wav"
    # 确保输出文件夹存在

    audio_folder = "SpeechTranslation/output_chunks/"
    os.makedirs(audio_folder, exist_ok=True)
    # 使用 AudioSegment 加载音频文件
    audio = AudioSegment.from_file(input_audio_file)
    audio_length_seconds = len(audio) / 1000.0
    # 将音频文件保存为 WAV 格式
    audio.export(output_wav_file, format="wav")
    asr = ASRExecutor()
    text=''
    if audio_length_seconds > 50:
        # 切分保存的 WAV 文件为多个50秒的片段
        # 定义切分的时长（单位：毫秒）
        chunk_length_ms = 50000  # 50秒

        # 使用 AudioSegment 再次加载 WAV 文件
        wav_audio = AudioSegment.from_file(output_wav_file)

        # 切分音频文件成多个片段
        chunks = make_chunks(wav_audio, chunk_length_ms)


        # 保存切分后的音频片段
        for i, chunk in enumerate(chunks):
            chunk = chunk.set_frame_rate(16000)
            output_file = f"{audio_folder}{i + 1}.wav"
            chunk.export(output_file, format="wav")

        print("音频文件已成功转换为 WAV 格式，并切分为多个50秒的片段，保存在:", audio_folder)

        audio_files = [os.path.join(audio_folder, file) for file in os.listdir(audio_folder) if file.endswith(".wav")]
        # 对音频文件列表按文件名排序
        audio_files.sort()
        for audio in audio_files:
            result = asr(audio_file=audio,device=device)
            text = result+text
    else:
        text = asr(audio_file=output_wav_file, device=device)
    # text_punc = TextExecutor()
    # translate_text = text_punc(text=text)
    model_path = 'models/translate/zh-en'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    pipeline = transformers.pipeline("translation", model=model, tokenizer=tokenizer)
    translate_text = pipeline(text)[0]['translation_text']
    # shutil.rmtree("exp")
    os.remove("SpeechTranslation/audio_recording.wav")
    os.remove("SpeechTranslation/output_audio.wav")
    return translate_text


SpeechTranslation/speech_translation_ch2en.py
tianclll/Ace-Translate
#! /usr/bin/env python
import os
import shutil
from paddlespeech.cli.st.infer import STExecutor
from pydub import AudioSegment
from pydub.utils import make_chunks
import configparser
config = configparser.ConfigParser()
config.read("config.conf")
device = config.get("devices", "device")
# 输入音频文件的路径（例如 MP3、FLAC、OGG 等）

def translate(input_audio_file):
    # 指定输出的 WAV 文件路径
    output_wav_file = "SpeechTranslation/output_audio.wav"

    # 使用 AudioSegment 加载音频文件
    audio = AudioSegment.from_file(input_audio_file)
    audio_length_seconds = len(audio) / 1000.0
    # 将音频文件保存为 WAV 格式
    audio.export(output_wav_file, format="wav")
    st = STExecutor()
    text=''
    if audio_length_seconds > 50:
        # 切分保存的 WAV 文件为多个50秒的片段
        # 定义切分的时长（单位：毫秒）
        chunk_length_ms = 50000  # 50秒

        # 使用 AudioSegment 再次加载 WAV 文件
        wav_audio = AudioSegment.from_file(output_wav_file)

        # 切分音频文件成多个片段
        chunks = make_chunks(wav_audio, chunk_length_ms)

        # 确保输出文件夹存在

        audio_folder = "SpeechTranslation/output_chunks/"
        os.makedirs(audio_folder, exist_ok=True)

        # 保存切分后的音频片段
        for i, chunk in enumerate(chunks):
            chunk = chunk.set_frame_rate(16000)
            output_file = f"{audio_folder}{i + 1}.wav"
            chunk.export(output_file, format="wav")

        print("音频文件已成功转换为 WAV 格式，并切分为多个50秒的片段，保存在:", audio_folder)

        audio_files = [os.path.join(audio_folder, file) for file in os.listdir(audio_folder) if file.endswith(".wav")]
        # 对音频文件列表按文件名排序
        audio_files.sort()
        for audio in audio_files:
            result = st(audio_file=audio,device=device)
            text = result+text
    else:
        text = st(audio_file=output_wav_file, device=device)
    shutil.rmtree("exp")
    os.remove("SpeechTranslation/audio_recording.wav")
    os.remove("SpeechTranslation/output_audio.wav")
    shutil.rmtree("SpeechTranslation/output_chunks")
    return text

SpeechTranslation/speech_translation_en2ch.py
tianclll/Ace-Translate
from datetime import timedelta
import math
import cv2
import os
from paddleocr import PaddleOCR
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from moviepy.editor import VideoFileClip
from moviepy.editor import TextClip
from moviepy.editor import CompositeVideoClip
running = True
def cvsecs(time):
    print(time)
    time_parts = time.split(':')
    time_parts[-1] = time_parts[-1].replace(',','.')
    seconds = time_parts[-1].replace(',','.')
    seconds = seconds.split('.')[0]
    milliseconds = "0"
    if len(time_parts[-1].split('.')) > 1:
        milliseconds = time_parts[-1].split('.')[1]
    milliseconds = milliseconds.replace(',', '.')
    hours = time_parts[0]
    minutes = time_parts[1]

    return float(hours) * 3600 + float(minutes) * 60 + float(seconds) + float(milliseconds) / 1000

def gen_video(video_path, output_file,fps):
    # 读取原始视频
    video = VideoFileClip(video_path)

    # 读取字幕文件
    subtitles = []
    with open("SubtitleTranslation/tmp.srt", "r") as srt_file:
        for line in srt_file:
            if not running:
                break
            line = line.strip()
            if line.isdigit():
                continue
            elif "-->" in line:
                start, end = line.split(" --> ")
            elif line:
                subtitles.append((start, end, line))
    print(subtitles)

    # 创建带有字幕的视频剪辑
    subtitled_video = video
    for subtitle in subtitles:
        if not running:
            break
        start_time = subtitle[0]
        end_time = subtitle[1]
        text = subtitle[2]

        # 将时间字符串转换为浮点数
        start_time = cvsecs(start_time)
        end_time = cvsecs(end_time)

        subtitle_clip = TextClip(text, fontsize=30, color='white', font='Arial', bg_color='black').set_position(
            ('center', 'bottom')).set_start(start_time).set_end(end_time)

        # 叠加字幕剪辑到视频上
        subtitled_video = CompositeVideoClip([subtitled_video.set_audio(None), subtitle_clip])


    # 设置视频的持续时间
    subtitled_video = subtitled_video.set_duration(video.duration)
    print(subtitled_video)
    # 添加原始视频的音频
    subtitled_video = subtitled_video.set_audio(video.audio)
    # 生成输出视频文件
    subtitled_video.write_videofile(output_file, codec='libx264',audio_codec="aac", fps=fps)

def translate(video_path, output_path, type):
    src_video = cv2.VideoCapture(video_path)
    ocr = PaddleOCR(use_angle_cls=False, lang='ch',cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/ch/ch_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/ch/ch_PP-OCRv4_rec_infer')
    model_path = 'models/translate/zh-en'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    pipeline = transformers.pipeline("translation", model=model, tokenizer=tokenizer)

    # 获取视频帧率
    fps = round(src_video.get(cv2.CAP_PROP_FPS))
    # 计算视频总帧数
    total_frame = int(src_video.get(cv2.CAP_PROP_FRAME_COUNT))

    # 设置字幕序号
    num = 0
    save_text = []
    save_time = []
    last_res = 'None'
    print(total_frame)

    if type == "输出字幕文件":
        video_name = os.path.basename(video_path)
        video_name = video_name.replace('mp4','srt')
        srt_path = os.path.join(output_path, video_name)
    else:
        srt_path = 'SubtitleTranslation/tmp.srt'

    with open(srt_path, 'w+') as f:
        for i in range(total_frame):
            if not running:
                break
            success, frame = src_video.read()
            # 在视频字幕提取任务中，为了提速改为了每秒抽一帧，因此所有时间都是整数
            if i % math.floor(fps) == 0:
                if success:
                    result = ocr.ocr(frame[-120:-30, :], cls=True)
                    if len(result[0]) > 0:
                        res = result[0][0][1][0]
                        res = pipeline(res)[0]['translation_text']
                        if ((res[:-1] not in last_res) and (last_res[:-1] not in res)):
                            # 检测到新字幕，录入并开始计时
                            last_res = res
                            # 更新序号
                            num += 1
                            # 更新开始时间
                            start_time = "%s,%03d" % (timedelta(seconds=i // fps), 0 * 1000)
                            # 一段语句的结束时间，和语速也有一定关系，读者可以自行调整
                            end_time = "%s,%03d" % (timedelta(seconds=i // fps), 1.5 * 1000)
                            f.write(str(num) + '\n')
                            f.write(start_time + ' --> ' + end_time + '\n')
                            f.write(res + '\n')
                            f.write('\n')
                            # 保存文稿信息，用于后续标点恢复和问答任务
                            save_text.append(res)
                            # 保存字幕 + 时间信息，用于快速截取视频片段
                            save_time.append(res + ' ' + str(i // fps))

    if type == "输出视频":
        output_path = os.path.join(output_path, os.path.basename(video_path))
        gen_video(video_path, output_path,fps)

    text = "已处理完成!"
    return text
def stop_translate():
    global running
    running = False
if __name__ == '__main__':
    translate('/Users/liuhongdi/Downloads/Trim.mp4', "/Users/liuhongdi/Downloads/", "输出视频")

SubtitleTranslation/subtitle_translation_ch2en.py
tianclll/Ace-Translate
from datetime import timedelta
import math
import cv2
import os
from paddleocr import PaddleOCR
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from moviepy.editor import VideoFileClip
from moviepy.editor import TextClip
from moviepy.editor import CompositeVideoClip
import sys
sys.path.append(os.getcwd())
import my_utils

running = True
def cvsecs(time):
    print(time)
    time_parts = time.split(':')
    time_parts[-1] = time_parts[-1].replace(',','.')
    seconds = time_parts[-1].replace(',','.')
    seconds = seconds.split('.')[0]
    milliseconds = "0"
    if len(time_parts[-1].split('.')) > 1:
        milliseconds = time_parts[-1].split('.')[1]
    milliseconds = milliseconds.replace(',', '.')
    hours = time_parts[0]
    minutes = time_parts[1]

    return float(hours) * 3600 + float(minutes) * 60 + float(seconds) + float(milliseconds) / 1000

def gen_video(video_path, output_file,fps):
    # 读取原始视频
    video = VideoFileClip(video_path)

    # 读取字幕文件
    subtitles = []
    with open("SubtitleTranslation/tmp.srt", "r") as srt_file:
        for line in srt_file:
            if not running:
                break
            line = line.strip()
            if line.isdigit():
                continue
            elif "-->" in line:
                start, end = line.split(" --> ")
            elif line:
                subtitles.append((start, end, line))
    print(subtitles)

    # 创建带有字幕的视频剪辑
    subtitled_video = video
    for subtitle in subtitles:
        if not running:
            break
        start_time = subtitle[0]
        end_time = subtitle[1]
        text = subtitle[2]

        # 将时间字符串转换为浮点数
        start_time = cvsecs(start_time)
        end_time = cvsecs(end_time)

        subtitle_clip = TextClip(text, fontsize=30, color='white', font='Arial', bg_color='black').set_position(
            ('center', 'bottom')).set_start(start_time).set_end(end_time)

        # 叠加字幕剪辑到视频上
        subtitled_video = CompositeVideoClip([subtitled_video.set_audio(None), subtitle_clip])

    # 设置视频的持续时间
    subtitled_video = subtitled_video.set_duration(video.duration)

    # 添加原始视频的音频
    subtitled_video = subtitled_video.set_audio(video.audio)

    # 生成输出视频文件
    subtitled_video.write_videofile(output_file, fps=fps, codec='libx264',audio_codec="aac")

def translate(video_path, output_path, type):
    src_video = cv2.VideoCapture(video_path)
    ocr = PaddleOCR(use_angle_cls=False, lang='en',cls_model_dir='models/ocr/cls/ch_ppocr_mobile_v2.0_cls_infer',det_model_dir='models/ocr/det/en/en_PP-OCRv4_det_infer',rec_model_dir='models/ocr/rec/en/en_PP-OCRv4_rec_infer')
    model_path = 'models/translate/en-zh'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    pipeline = transformers.pipeline("translation", model=model, tokenizer=tokenizer)

    # 获取视频帧率
    fps = round(src_video.get(cv2.CAP_PROP_FPS))
    # 计算视频总帧数
    total_frame = int(src_video.get(cv2.CAP_PROP_FRAME_COUNT))

    # 设置字幕序号
    num = 0
    save_text = []
    save_time = []
    last_res = 'None'
    print(total_frame)

    if type == "输出字幕文件":
        video_name = os.path.basename(video_path)
        video_name = video_name.replace('mp4','srt')
        srt_path = os.path.join(output_path, video_name)
    else:
        srt_path = 'SubtitleTranslation/tmp.srt'

    with open(srt_path, 'w+') as f:
        for i in range(total_frame):
            if not running:
                break
            success, frame = src_video.read()

            if i % math.floor(fps) == 0:
                if success:
                    result = ocr.ocr(frame[-120:-30, :], cls=True)
                    if len(result[0]) > 0:
                        res = result[0][0][1][0]
                        res = pipeline(res)[0]['translation_text']
                        res = my_utils.do_sentence(res)
                        if ((res[:-1] not in last_res) and (last_res[:-1] not in res)):
                            # 检测到新字幕，录入并开始计时
                            last_res = res
                            # 更新序号
                            num += 1
                            # 更新开始时间
                            start_time = "%s,%03d" % (timedelta(seconds=i // fps), 0 * 1000)
                            # 一段语句的结束时间，和语速也有一定关系，读者可以自行调整
                            end_time = "%s,%03d" % (timedelta(seconds=i // fps), 1.5 * 1000)
                            f.write(str(num) + '\n')
                            f.write(start_time + ' --> ' + end_time + '\n')
                            f.write(res + '\n')
                            f.write('\n')
                            # 保存文稿信息，用于后续标点恢复和问答任务
                            save_text.append(res)
                            # 保存字幕 + 时间信息，用于快速截取视频片段
                            save_time.append(res + ' ' + str(i // fps))

    if type == "输出视频":
        output_path = os.path.join(output_path, os.path.basename(video_path))
        gen_video(video_path, output_path,fps)

    text = "已处理完成!"
    return text
def stop_translate():
    global running
    running = False
if __name__ == '__main__':
    translate('/Users/liuhongdi/Downloads/Trim.mp4', "/Users/liuhongdi/Downloads/", "输出视频")

SubtitleTranslation/subtitle_translation_en2ch.py
tianclll/Ace-Translate
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
def translate(text):
    model_path = 'models/translate/zh-en'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    translate_model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    pipeline = transformers.pipeline("translation", model=translate_model, tokenizer=tokenizer)
    translate_text = pipeline(text)[0]['translation_text']
    return translate_text
Translation/translation_ch2en.py
tianclll/Ace-Translate

import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import sys
import os
sys.path.append(os.getcwd())
import my_utils
def translate(text):
    model_path = 'models/translate/en-zh'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    translate_model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    pipeline = transformers.pipeline("translation", model=translate_model, tokenizer=tokenizer)
    translate_text = pipeline(text)[0]['translation_text']
    translate_text = my_utils.do_sentence(translate_text)
    return translate_text
Translation/translation_en2ch.py
tianclll/Ace-Translate
import glob
import os
import shutil
import wave
import pystray
import pyaudio
import threading
import subprocess
import tkinter as tk
from PIL import Image
from tkinter import ttk
from tkinter import filedialog, scrolledtext
from ttkthemes import ThemedStyle  # 导入 ThemedStyle


class TranslationApp:
    def __init__(self, root):
        self.root = root
        self.root.title("全能翻译")
        # 创建 ttk 样式dff
        self.style = ThemedStyle(root)
        self.style.set_theme("radiance")  # 设置 Radiance 主题
        self.style.configure("Selected.TLabel", background="#F6F4F3")  # 选中模块的标签样式
        self.style.configure("Unselected.TLabel", background="lightgray")  # 未选中模块的标签样式

        self.sidebar = ttk.Frame(root, width=200)
        self.sidebar.pack(side="left", fill="y",expand=True)

        self.modules = ["文本翻译","划词翻译", "截图翻译", "语音翻译", "视频翻译", "文件翻译", "文档图片翻译"]
        self.module_labels = []
        for module in self.modules:
            label_frame = tk.Frame(self.sidebar)
            label_frame.pack(fill="x")

            label_style = "Selected.TLabel" if module == self.modules[0] else "Unselected.TLabel"  # 设置不同模块的标签样式
            label = ttk.Label(label_frame, text=module, font=("Helvetica", 12), style=label_style)
            label.pack(fill="x", padx=10, pady=5)
            label.bind("<Button-1>", lambda event, m=module: self.run_module(m))
            self.module_labels.append(label)

        self.content = tk.Frame(root, padx=20, pady=20)
        self.content.pack(side="right", fill="both", expand=True)
        self.content.configure(bg="#F6F4F3")  # 设置背景颜色

        self.recording_thread = None  # 用于存储录音线程
        self.audio_file_path = None  # 存储录音文件的路径
        self.selected_module = None  # 保存当前选中的模块
        self.threads = {}  # 存储模块对应的线程
        self.translation_thread = None     # 创建一个用于翻译的线程
        self.voice_translation_thread = None  # 创建一个用于语音翻译的线程
        self.video_translation_thread = None
        self.run_module(self.modules[0])  # 默认显示第一个模块
        self.running_processes = {}  # 用于保存各模块的进程对象
        self.create_tray_icon()
        self.root.protocol('WM_DELETE_WINDOW', self.hide_window)
    def create_tray_icon(self):
        icon_path = "static/translate.png"  # 修改为你的图标路径

        # 读取图标并创建菜单项
        icon_image = Image.open(icon_path)
        menu = (
            pystray.MenuItem("打开", self.show_window),
            pystray.MenuItem("退出", self.quit_app)
        )

        # 创建系统托盘图标
        self.tray_icon = pystray.Icon("app_name", icon_image, "Ace Translate", menu)
        threading.Thread(target=self.tray_icon.run, daemon=True).start()
    def show_window(self):
        if not self.root.winfo_viewable():
            self.root.deiconify()  # 显示窗口
            self.root.lift()  # 将窗口置于顶层
        else:
            self.root.iconify()  # 隐藏窗口
    def hide_window(self):
        self.root.withdraw()
    def run_module(self, module):
        if module in self.threads and self.threads[module].is_alive():
            return  # 如果模块的线程正在运行，则不做任何操作
        self.selected_module = module
        self.update_sidebar_colors()
        self.clear_content()
        if module == "文本翻译":
            thread = threading.Thread(target=self.run_translation_module)
        elif module == "划词翻译":
            thread = threading.Thread(target=self.run_hyper_translation_module)
        elif module == "截图翻译":
            thread = threading.Thread(target=self.run_screenshot_module)
        elif module == "语音翻译":
            thread = threading.Thread(target=self.run_voice_translation_module)
        elif module == "视频翻译":
            thread = threading.Thread(target=self.run_video_module)
        elif module == "文件翻译":
            thread = threading.Thread(target=self.run_pdf_module)
        elif module == "文档图片翻译":
            thread = threading.Thread(target=self.run_image_module)

        self.threads[module] = thread
        thread.start()

    def clear_content(self):
        for widget in self.content.winfo_children():
            widget.destroy()

    def update_sidebar_colors(self):
        for label in self.module_labels:
            label_frame = label.master
            if label.cget("text") == self.selected_module:
                label.config(style="Selected.TLabel")  # 使用 Selected.TLabel 样式
                label_frame.config(bg="#F6F4F3")
            else:
                label.config(style="Unselected.TLabel")  # 使用 Unselected.TLabel 样式
                label_frame.config(bg="lightgray")

    def stop_module(self, module_name):
        if module_name == "视频翻译":
            from SubtitleTranslation import subtitle_translation_ch2en,subtitle_translation_en2ch
            subtitle_translation_ch2en.stop_translate()
            subtitle_translation_en2ch.stop_translate()
        if module_name in self.running_processes:
            process = self.running_processes[module_name]
            if os.name == 'nt':  # Windows系统
                process.terminate()
            else:  # Linux或类Unix系统
                process.kill()
            del self.running_processes[module_name]
    def run_translation_module(self):

        translation_module = ttk.Frame(self.content)
        label = ttk.Label(translation_module, text="文本翻译", font=("Helvetica", 16))
        label.pack()

        # 添加单选框，默认选中"汉译英"
        self.translation_direction = tk.StringVar(value="汉译英")
        direction_frame = ttk.Frame(translation_module)
        direction_frame.pack()

        to_english_radio = ttk.Radiobutton(direction_frame, text="汉译英", variable=self.translation_direction,
                                           value="汉译英")
        to_english_radio.pack(side="left")

        to_chinese_radio = ttk.Radiobutton(direction_frame, text="英译汉", variable=self.translation_direction,
                                           value="英译汉")
        to_chinese_radio.pack(side="left")

        # 创建文本输入框，将其添加到 translation_module
        entry = ttk.Entry(translation_module, width=20)
        entry.pack(padx=10, pady=10)

        # 创建翻译按钮，使用 lambda 包装函数
        translate_button = ttk.Button(translation_module, text="翻译",
                                      command=lambda: self.start_translation_module(entry.get()))
        translate_button.pack(pady=10)

        # 创建文本标签用于显示翻译结果，将其添加到 translation_module
        self.translate_translated_text = tk.StringVar()
        result_entry = ttk.Entry(translation_module, textvariable=self.translate_translated_text, width=20)
        result_entry.configure(state="readonly")
        result_entry.pack(padx=10, pady=10)

        # 添加 translation_module 到 content
        translation_module.pack(fill="both", expand=True)

    def start_translation_module(self, translate_text):
        if self.translation_thread is not None and self.translation_thread.is_alive():
            return

        self.translation_thread = threading.Thread(target=self.run_translation, args=(translate_text,))
        self.translation_thread.start()


    def run_translation(self,translate_text):
        from Translation import translation_ch2en, translation_en2ch
        translation_direction = self.translation_direction.get()  # 获取选择的翻译方向
        if translation_direction == "英译汉":
            text = translation_en2ch.translate(translate_text)
        else:
            text = translation_ch2en.translate(translate_text)
        self.translate_translated_text.set(text)
    def run_hyper_translation_module(self):
        module_name = "划词翻译"

        translation_module = ttk.Frame(self.content)

        label = ttk.Label(translation_module, text="划词翻译",font=("Helvetica", 16))
        label.pack()
        # 添加单选框，默认选中"汉译英"
        self.translation_direction = tk.StringVar(value="汉译英")
        direction_frame = ttk.Frame(translation_module)
        direction_frame.pack()

        to_english_radio = ttk.Radiobutton(direction_frame, text="汉译英", variable=self.translation_direction, value="汉译英")
        to_english_radio.pack(side="left")

        to_chinese_radio = ttk.Radiobutton(direction_frame, text="英译汉", variable=self.translation_direction, value="英译汉")
        to_chinese_radio.pack(side="left")
        start_button = ttk.Button(translation_module, text="开始", command=lambda: self.start_hyper_translation_module(module_name))
        start_button.pack()

        stop_button = ttk.Button(translation_module, text="结束", command=lambda: self.stop_module(module_name))
        stop_button.pack()

        translation_module.pack(fill="both", expand=True)

    def start_hyper_translation_module(self, module_name):
        if module_name not in self.running_processes:
            translation_direction = self.translation_direction.get()  # 获取选择的翻译方向
            if translation_direction == "英译汉":
                process = subprocess.Popen(['python', 'HyperTranslation/hyper_translation_en2ch.py',translation_direction])
            else:
                process = subprocess.Popen(['python', 'HyperTranslation/hyper_translation_ch2en.py', translation_direction])
            self.running_processes[module_name] = process

    # 其他模块函数省略...
    def run_screenshot_module(self):
        module_name = "截图翻译"
        screenshot_module = ttk.Frame(self.content)
        label = ttk.Label(screenshot_module, text="截图翻译", font=("Helvetica", 16))
        label.pack()
        # 添加单选框，默认选中"汉译英"
        self.translation_direction = tk.StringVar(value="汉译英")
        direction_frame = ttk.Frame(screenshot_module)
        direction_frame.pack()

        to_english_radio = ttk.Radiobutton(direction_frame, text="汉译英", variable=self.translation_direction, value="汉译英",)
        to_english_radio.pack(side="left")

        to_chinese_radio = ttk.Radiobutton(direction_frame, text="英译汉", variable=self.translation_direction, value="英译汉")
        to_chinese_radio.pack(side="left")
        start_button = ttk.Button(screenshot_module, text="开始", command=lambda: self.start_screenshot_module(module_name))
        start_button.pack()

        stop_button = ttk.Button(screenshot_module, text="结束", command=lambda: self.stop_module(module_name))
        stop_button.pack()

        screenshot_module.pack(fill="both", expand=True)
    def start_screenshot_module(self, module_name):
        if module_name not in self.running_processes:
            translation_direction = self.translation_direction.get()  # 获取选择的翻译方向
            if translation_direction == "英译汉":
                process = subprocess.Popen(['python', 'ScreenshotTranslation/screenshot_translation_en2ch.py'])
            else:
                process = subprocess.Popen(['python', 'ScreenshotTranslation/screenshot_translation_ch2en.py'])
            self.running_processes[module_name] = process

    def run_voice_translation_module(self):
        self.voice_translation_module = ttk.Frame(self.content)
        # self.voice_translation_module.pack_forget()  # 隐藏模块

        label = ttk.Label(self.voice_translation_module, text="语音翻译", font=("Helvetica", 16))
        label.pack()

        self.translation_direction = tk.StringVar(value="英译汉")
        direction_frame = ttk.Frame(self.voice_translation_module)
        direction_frame.pack()

        to_english_radio = ttk.Radiobutton(direction_frame, text="英译汉", variable=self.translation_direction,
                                           value="英译汉")
        to_english_radio.pack(side="left")

        to_chinese_radio = ttk.Radiobutton(direction_frame, text="汉译英", variable=self.translation_direction,
                                           value="汉译英")
        to_chinese_radio.pack(side="left")

        # 使用ttk.LabelFrame来包裹按钮
        button_frame = ttk.LabelFrame(self.voice_translation_module, text="录音控制")
        button_frame.pack()

        # 启动录音按钮
        self.record_audio_button = ttk.Button(button_frame, text="开始录音", command=self.start_record)
        self.record_audio_button.grid(row=0, column=0, padx=10, pady=5)

        # 停止录音按钮
        self.stop_button = ttk.Button(button_frame, text="停止录音", command=self.stop_record, state=tk.DISABLED)
        self.stop_button.grid(row=0, column=1, padx=10, pady=5)

        # 上传音频文件按钮
        self.upload_audio_button = ttk.Button(button_frame, text="上传音频文件", command=self.upload_audio_file)
        self.upload_audio_button.grid(row=1, column=0, padx=10, pady=5)

        # 开始翻译按钮
        self.start_translation_button = ttk.Button(button_frame, text="开始翻译",
                                                   command=lambda: self.start_voice_translation("语音翻译"),
                                                   state=tk.DISABLED)
        self.start_translation_button.grid(row=1, column=1, padx=10, pady=5)

        # 翻译结果文本框
        self.voice_translation_result = scrolledtext.ScrolledText(self.voice_translation_module, padx=10, pady=10,
                                                                  height=6, width=40)
        self.voice_translation_result.pack()

        self.voice_translation_module.pack(fill="both", expand=True)

    def start_voice_translation(self,module_name):

        if self.voice_translation_thread is not None and self.voice_translation_thread.is_alive():
            return

        self.voice_translation_thread = threading.Thread(target=self.run_voice_translation, args=(module_name,))
        self.voice_translation_thread.start()

    is_recording = False  # 添加一个标志来控制录音状态
    def run_voice_translation(self, module_name):
        # 这个函数包含了你之前在 start_voice_translation_module 函数中的语音翻译逻辑
        from SpeechTranslation import speech_translation_ch2en, speech_translation_en2ch

        # 获取用户选择的翻译方向
        translation_direction = self.translation_direction.get()

        if translation_direction == "英译汉":
            text = speech_translation_en2ch.translate(self.audio_file_path)
        else:
            text = speech_translation_ch2en.translate(self.audio_file_path)

        # 在文本框中显示翻译文本
        self.voice_translation_result.delete("1.0", "end")  # 清空文本框
        self.voice_translation_result.insert("end", str(text))  # 插入翻译文本

    def start_record(self):
        # 启用停止录音按钮
        self.stop_button.config(state=tk.NORMAL)
        # 禁用开始录音按钮
        self.record_audio_button.config(state=tk.DISABLED)
        # 禁用上传音频文件按钮
        self.upload_audio_button.config(state=tk.DISABLED)

        self.audio_thread = threading.Thread(target=self.record_audio)
        self.audio_thread.start()

    def record_audio(self):
        CHUNK = 1024
        FORMAT = pyaudio.paInt16
        CHANNELS = 1
        RATE = 16000
        audio = pyaudio.PyAudio()

        stream = audio.open(format=FORMAT, channels=CHANNELS,
                            rate=RATE, input=True,
                            frames_per_buffer=CHUNK)
        frames = []

        print("开始录音...")
        global is_recording
        is_recording = True
        while is_recording:
            data = stream.read(CHUNK)
            frames.append(data)
        print("录音结束...")

        stream.stop_stream()
        stream.close()
        audio.terminate()

        # 保存录音文件
        self.audio_file_path = "SpeechTranslation/audio_recording.wav"
        with wave.open(self.audio_file_path, 'wb') as wf:
            wf.setnchannels(CHANNELS)
            wf.setsampwidth(audio.get_sample_size(FORMAT))
            wf.setframerate(RATE)
            wf.writeframes(b''.join(frames))

        # 启用开始录音按钮
        self.record_audio_button.config(state=tk.NORMAL)
        # 启用上传音频文件按钮
        self.upload_audio_button.config(state=tk.NORMAL)
        # 禁用停止录音按钮
        self.stop_button.config(state=tk.DISABLED)
        # 启用开始翻译按钮
        self.start_translation_button.config(state=tk.NORMAL)
    def upload_audio_file(self):
        file_path = filedialog.askopenfilename(filetypes=[("Audio Files", "*.wav *.mp3")])
        if file_path:
            self.audio_file_path = file_path
            self.start_translation_button.config(state=tk.NORMAL)
    def stop_record(self):
        global is_recording  # 使用全局标志来控制录音状态
        is_recording = False  # 停止录音
        # 禁用停止录音按钮
        self.stop_button.config(state=tk.DISABLED)
        # 启用开始录音按钮
        self.record_audio_button.config(state=tk.NORMAL)
        # 启用上传音频文件按钮
        self.upload_audio_button.config(state=tk.NORMAL)
        # 禁用开始翻译按钮
        self.start_translation_button.config(state=tk.DISABLED)
    def run_video_module(self):
        module_name = "视频翻译"
        video_module = ttk.Frame(self.content)
        label = ttk.Label(video_module, text="视频翻译", font=("Helvetica", 16))
        label.grid(row=0, column=0, columnspan=2)

        # 添加单选框，默认选中"汉译英"
        self.translation_direction = tk.StringVar(value="汉译英")
        direction_frame = ttk.Frame(video_module)
        direction_frame.grid(row=1, column=0, columnspan=2, pady=(10, 0))

        to_english_radio = ttk.Radiobutton(direction_frame, text="汉译英", variable=self.translation_direction,
                                           value="汉译英")
        to_english_radio.grid(row=0, column=0)

        to_chinese_radio = ttk.Radiobutton(direction_frame, text="英译汉", variable=self.translation_direction,
                                           value="英译汉")
        to_chinese_radio.grid(row=0, column=1)

        # upload_button = ttk.Button(video_module, text="上传文件", command=lambda: self.upload_pdf(module_name))
        # upload_button.grid(row=3, column=0, columnspan=2, pady=(10, 0))

        # 添加单选框，默认选中"汉译英"
        self.output_type = tk.StringVar(value="输出字幕文件")
        type = ttk.Frame(video_module)
        type.grid(row=2, column=0, columnspan=2, pady=(10, 0))

        to_srt = ttk.Radiobutton(type, text="输出字幕文件", variable=self.output_type,
                                           value="输出字幕文件")
        to_srt.grid(row=0, column=0)

        to_video = ttk.Radiobutton(type, text="输出视频", variable=self.output_type,
                                           value="输出视频")
        to_video.grid(row=0, column=1)
        upload_button = ttk.Button(video_module, text="上传视频文件", command=lambda: self.upload_video(module_name))
        upload_button.grid(row=3, column=0, columnspan=2, pady=(10, 0))


        self.save_path = tk.StringVar()
        save_path_entry = ttk.Entry(video_module, textvariable=self.save_path)
        save_path_entry.grid(row=4, column=0, pady=(10, 0))

        choose_path_button = ttk.Button(video_module, text="选择输出路径", command=self.choose_video_save_path)
        choose_path_button.grid(row=4, column=1, pady=(10, 0))

        start_button = ttk.Button(video_module, text="开始", command=lambda: self.start_video_module(module_name))
        start_button.grid(row=5, column=0, pady=(10, 0),sticky="w")

        stop_button = ttk.Button(video_module, text="结束", command=lambda: self.stop_module(module_name))
        stop_button.grid(row=5, column=1, pady=(10, 0),sticky="w")

        video_module.grid(row=0, column=0, sticky="nsew")

        self.video_translation_result = scrolledtext.ScrolledText(video_module, padx=10, pady=10, height=2, width=40)
        self.video_translation_result.grid(row=6, column=0, columnspan=2, pady=(10, 0))

        self.video_translation_result.grid(row=6, column=0, columnspan=2, pady=(10, 0))

        self.video_module = video_module
    def upload_video(self, module_name):
        file_path = filedialog.askopenfilename(filetypes=[("Video Files", "*.mp4")])
        if file_path:
            self.video_file_path = file_path
    def choose_video_save_path(self):
        chosen_path = filedialog.askdirectory()
        if chosen_path:
            self.save_path.set(chosen_path)
            # self.save_path = chosen_path
            # save_path_entry = ttk.Entry(self.video_module, textvariable=self.save_path)
            # save_path_entry.grid(row=3, column=0, pady=(10, 0))
    def start_video_module(self,module_name):
        if self.video_translation_thread is not None and self.video_translation_thread.is_alive():
            return
        self.video_translation_thread = threading.Thread(target=self.run_video_translation)
        self.video_translation_thread.start()

    def run_video_translation(self):
        # 这个函数包含了你之前在 start_voice_translation_module 函数中的语音翻译逻辑
        from SubtitleTranslation import subtitle_translation_ch2en, subtitle_translation_en2ch
        self.video_translation_result.delete("1.0", "end")  # 清空文本框
        self.video_translation_result.insert("end", "开始运行...\n")  # 插入翻译文本
        # 获取用户选择的翻译方向
        translation_direction = self.translation_direction.get()
        output_type = self.output_type.get()
        if translation_direction == "英译汉":
            text = subtitle_translation_en2ch.translate(self.video_file_path, self.save_path.get(), output_type)
        else:
            text = subtitle_translation_ch2en.translate(self.video_file_path, self.save_path.get(), output_type)

        # 在文本框中显示翻译文本
        # self.video_translation_result.delete("1.0", "end")  # 清空文本框
        self.video_translation_result.insert("end", str(text))  # 插入翻译文本
    def start_pdf_module(self,module_name):
        if module_name not in self.running_processes:
            log_message = f"开始运行 {module_name} 模块...\n将存入{self.save_path.get()}中."
            self.log_text.insert("end", log_message)
            self.log_text.see("end")  # 自动滚动到底部
            file_extension = self.pdf_file_path.split(".")[-1]
            if file_extension in ['xls','xlsx']:
                translation_direction = self.translation_direction.get()  # 获取选择的翻译方向
                if translation_direction == "英译汉":
                    process = subprocess.Popen(
                        ['python', 'FileTranslation/Excel/excel_translation_en2ch.py', '{}'.format(self.pdf_file_path),
                         '{}'.format(self.save_path.get())])
                else:
                    process = subprocess.Popen(
                        ['python', 'FileTranslation/Excel/excel_translation_ch2en.py', '{}'.format(self.pdf_file_path),
                         '{}'.format(self.save_path.get())])
                self.running_processes[module_name] = process
            elif file_extension == 'pdf':
                translation_direction = self.translation_direction.get()  # 获取选择的翻译方向
                if translation_direction == "英译汉":
                    process = subprocess.Popen(['python', 'FileTranslation/PDF/pdf_translation_en2ch.py','{}'.format(self.pdf_file_path),'{}'.format(self.save_path.get())])
                else:
                    process = subprocess.Popen(['python', 'FileTranslation/PDF/pdf_translation_ch2en.py', '{}'.format(self.pdf_file_path),'{}'.format(self.save_path.get())])
                self.running_processes[module_name] = process
            elif file_extension == 'txt':
                translation_direction = self.translation_direction.get()  # 获取选择的翻译方向
                if translation_direction == "英译汉":
                    process = subprocess.Popen(['python', 'FileTranslation/TXT/txt_translation_en2ch.py','{}'.format(self.pdf_file_path),'{}'.format(self.save_path.get())])
                else:
                    process = subprocess.Popen(['python', 'FileTranslation/TXT/txt_translation_ch2en.py', '{}'.format(self.pdf_file_path),'{}'.format(self.save_path.get())])
                self.running_processes[module_name] = process
            elif file_extension == 'docx':
                translation_direction = self.translation_direction.get()  # 获取选择的翻译方向
                if translation_direction == "英译汉":
                    process = subprocess.Popen(
                        ['python', 'FileTranslation/Word/word_translation_en2ch.py', '{}'.format(self.pdf_file_path),
                         '{}'.format(self.save_path.get())])
                else:
                    process = subprocess.Popen(
                        ['python', 'FileTranslation/Word/word_translation_ch2en.py', '{}'.format(self.pdf_file_path),
                         '{}'.format(self.save_path.get())])
                self.running_processes[module_name] = process
            elif file_extension == 'pptx':
                translation_direction = self.translation_direction.get()  # 获取选择的翻译方向
                if translation_direction == "英译汉":
                    process = subprocess.Popen(
                        ['python', 'FileTranslation/PPT/ppt_translation_en2ch.py', '{}'.format(self.pdf_file_path),
                         '{}'.format(self.save_path.get())])
                else:
                    process = subprocess.Popen(
                        ['python', 'FileTranslation/PPT/ppt_translation_ch2en.py', '{}'.format(self.pdf_file_path),
                         '{}'.format(self.save_path.get())])
                self.running_processes[module_name] = process
            else:
                translation_direction = self.translation_direction.get()  # 获取选择的翻译方向
                if translation_direction == "英译汉":
                    process = subprocess.Popen(
                        ['python', 'FileTranslation/Image/image_translation_en2ch.py', '{}'.format(self.pdf_file_path),
                         '{}'.format(self.save_path.get())])
                else:
                    process = subprocess.Popen(
                        ['python', 'FileTranslation/Image/image_translation_ch2en.py', '{}'.format(self.pdf_file_path),
                         '{}'.format(self.save_path.get())])
                self.running_processes[module_name] = process
    def run_pdf_module(self):
        module_name = "文件翻译"
        pdf_module = ttk.Frame(self.content)
        label = ttk.Label(pdf_module, text="文件翻译", font=("Helvetica", 16))
        label.grid(row=0, column=0, columnspan=2)

        # 添加单选框，默认选中"汉译英"
        self.translation_direction = tk.StringVar(value="汉译英")
        direction_frame = ttk.Frame(pdf_module)
        direction_frame.grid(row=1, column=0, columnspan=2, pady=(10, 0))

        to_english_radio = ttk.Radiobutton(direction_frame, text="汉译英", variable=self.translation_direction,
                                           value="汉译英")
        to_english_radio.grid(row=0, column=0)

        to_chinese_radio = ttk.Radiobutton(direction_frame, text="英译汉", variable=self.translation_direction,
                                           value="英译汉")
        to_chinese_radio.grid(row=0, column=1)

        upload_button = ttk.Button(pdf_module, text="上传文件", command=lambda: self.upload_pdf(module_name))
        upload_button.grid(row=2, column=0, columnspan=2, pady=(10, 0))


        self.save_path = tk.StringVar()
        save_path_entry = ttk.Entry(pdf_module, textvariable=self.save_path)
        save_path_entry.grid(row=3, column=0, pady=(10, 0))

        choose_path_button = ttk.Button(pdf_module, text="选择输出路径", command=self.choose_pdf_save_path)
        choose_path_button.grid(row=3, column=1, pady=(10, 0))

        start_button = ttk.Button(pdf_module, text="开始", command=lambda: self.start_pdf_module(module_name))
        start_button.grid(row=4, column=0, pady=(10, 0),sticky="w")

        stop_button = ttk.Button(pdf_module, text="结束", command=lambda: self.stop_module(module_name))
        stop_button.grid(row=4, column=1, pady=(10, 0),sticky="w")

        pdf_module.grid(row=0, column=0, sticky="nsew")

        self.log_text = scrolledtext.ScrolledText(pdf_module, padx=10, pady=10, height=2, width=40)
        self.log_text.grid(row=5, column=0, columnspan=2, pady=(10, 0))

        self.log_text.grid(row=5, column=0, columnspan=2, pady=(10, 0))

        self.pdf_module = pdf_module  # 存储 PDF 翻译模块的 Frame

    def choose_pdf_save_path(self):
        chosen_path = filedialog.askdirectory()
        if chosen_path:
            self.save_path.set(chosen_path)
            # self.save_path = chosen_path
            save_path_entry = ttk.Entry(self.pdf_module, textvariable=self.save_path)
            save_path_entry.grid(row=3, column=0, pady=(10, 0))
    def choose_img_save_path(self):
        chosen_path = filedialog.askdirectory()
        if chosen_path:
            self.save_path.set(chosen_path)
            # self.save_path = chosen_path
            save_path_entry = ttk.Entry(self.image_module, textvariable=self.save_path)
            save_path_entry.grid(row=3, column=0, pady=(10, 0))
    def upload_pdf(self, module_name):
        filetypes = [
            ("PDF Files", "*.pdf"),
            ("Image Files", ("*.png", "*.jpg", "*.jpeg")),
            ("Excel Files", ("*.xls", "*.xlsx")),
            ("Text Files", "*.txt"),
            ("Word Files","*.docx"),
            ("PowerPoint Files","*.pptx")
        ]
        file_path = filedialog.askopenfilename(filetypes=filetypes)
        if file_path:
            self.pdf_file_path = file_path
    def start_image_module(self,module_name):
        if module_name not in self.running_processes:
            log_message = f"开始运行 {module_name} 模块...\n将存入{self.save_path.get()}中."
            self.log_text.insert("end", log_message)
            self.log_text.see("end")  # 自动滚动到底部
            translation_direction = self.translation_direction.get()  # 获取选择的翻译方向
            if translation_direction == "英译汉":
                process = subprocess.Popen(['python','DocImgTranslation/ocr_image_en2ch.py','-i',f'{self.pdf_file_path}','-o',f'{self.save_path.get()}'])
            else:

                process = subprocess.Popen(['python','DocImgTranslation/ocr_image_ch2en.py','-i',f'{self.pdf_file_path}','-o',f'{self.save_path.get()}'])
            self.running_processes[module_name] = process


    def run_image_module(self):

        module_name = "文档图片翻译"
        image_module = ttk.Frame(self.content)
        label = ttk.Label(image_module, text="文档图片翻译", font=("Helvetica", 16))
        label.grid(row=0, column=0, columnspan=2)

        direction_frame = ttk.Frame(image_module)
        direction_frame.grid(row=1, column=0, columnspan=2, pady=(10, 0))

        to_english_radio = ttk.Radiobutton(direction_frame, text="汉译英", variable=self.translation_direction,
                                           value="汉译英")
        to_english_radio.pack(side="left")

        to_chinese_radio = ttk.Radiobutton(direction_frame, text="英译汉", variable=self.translation_direction,
                                           value="英译汉")
        to_chinese_radio.pack(side="left")

        upload_button = ttk.Button(image_module, text="上传图片文件", command=lambda: self.upload_image(module_name))
        upload_button.grid(row=2, column=0, columnspan=2, pady=(10, 0))

        self.save_path = tk.StringVar()
        save_path_entry = ttk.Entry(image_module, textvariable=self.save_path)
        save_path_entry.grid(row=3, column=0, pady=(10, 0), sticky="ew")

        choose_path_button = ttk.Button(image_module, text="选择输出路径", command=self.choose_img_save_path)
        choose_path_button.grid(row=3, column=1, pady=(10, 0), sticky="w")

        start_button = ttk.Button(image_module, text="开始", command=lambda: self.start_image_module(module_name))
        start_button.grid(row=4, column=0, pady=(10, 0), sticky="w")

        stop_button = ttk.Button(image_module, text="结束", command=lambda: self.stop_module(module_name))
        stop_button.grid(row=4, column=1, pady=(10, 0), sticky="w")
        image_module.grid(row=0, column=0, sticky="nsew")
        self.log_text = scrolledtext.ScrolledText(image_module, padx=10, pady=10, height=2, width=40)
        self.log_text.grid(row=5, column=0, columnspan=2, pady=(10, 0))
        self.log_text.grid(row=5, column=0, columnspan=2, pady=(10, 0))
        image_module.pack(fill="both", expand=True)
        self.image_module = image_module
    def upload_image(self, module_name):
        file_path = filedialog.askopenfilename(filetypes=[("image Files", "*.jpg *.png")])
        if file_path:
            self.pdf_file_path = file_path
    def quit_app(self):
        if os.path.exists("exp"):
            shutil.rmtree("exp")
        mp3_pattern = "*.mp3"
        # 使用 glob 模块查找匹配的文件
        mp3_files = glob.glob(mp3_pattern)
        if mp3_files:
            for mp3_file in mp3_files:
                os.remove(mp3_file)
        # 终止所有模块线程
        for module, thread in self.threads.items():
            if thread.is_alive():
                thread.join()
        for i in self.running_processes.keys():
            process = self.running_processes[i]
            if os.name == 'nt':  # Windows系统
                process.terminate()
            else:  # Linux或类Unix系统
                process.kill()
        self.root.quit()
if __name__ == '__main__':

    root = tk.Tk()
    app = TranslationApp(root)
    root.mainloop()

main.py
tianclll/Ace-Translate
import re
from nltk import word_tokenize
import numpy as np
import jieba

def process_word(text):

    if re.search(r'[\u4e00-\u9fff]', text):  # 判断是否包含中文字符
        text = text.replace(' ', '')
        if len(text) > 1:
            result = re.sub(r'([\u4e00-\u9fff])\1+', r'\1', text)  # 替换连续的相同中文字为第一个中文字
            return result
        else:
            return text
    else:
        result = re.sub(r'([a-zA-Z])\1{2,}', r'\1\1', text)  # 替换连续的相同字母为第一个字母
        return result
def remove_last_punctuation(text):
    # 匹配最后一个标点符号的正则表达式
    pattern = r'[,.?;]$'
    punctuation_list = []
    cleaned_list = []
    for sentence in text:
        match = re.search(pattern, sentence)
        if match:
            # 如果匹配到了标点符号，则去除标点符号并记录到结果列表中
            cleaned_sentence = re.sub(pattern, '', sentence)
            punctuation = match.group()
        else:
            # 如果没有匹配到标点符号，则保持原始文本并将结果列表中添加空字符串
            cleaned_sentence = sentence
            punctuation = ''
        punctuation_list.append(punctuation)
        cleaned_list.append(cleaned_sentence)
    return punctuation_list, cleaned_list

def cosine_similarity(u, v):
    dot_product = np.dot(u, v)
    norm_u = np.linalg.norm(u)
    norm_v = np.linalg.norm(v)
    similarity = dot_product / (norm_u * norm_v)
    return similarity
def do_sentence(translated_text):
    if translated_text == '':
        return ''
    translated_text = process_word(translated_text)
    print(translated_text)
    tokenized_sentence = list(jieba.cut(translated_text))
    # 保留第一个词语
    words = [tokenized_sentence[0]]  # 保留第一个词语
    for i in range(1, len(tokenized_sentence)):
        if tokenized_sentence[i] != tokenized_sentence[i - 1]:
            words.append(tokenized_sentence[i])
    words = [process_word(word) for word in words]
    print(words)
    filtered_sentence = "".join(words)
    sentences = []
    j = ''
    for i in words:
        j = j + i
        if i[-1] in [',', '?', '!', ')', ';', ' ', '...']:
            sentences.append(j)
            j = ''
    sentences.append(j)
    if len(sentences) <= 1:
        return filtered_sentence
    # 分词和构建词袋表示
    punctuation_list, sentences = remove_last_punctuation(sentences)
    print(punctuation_list)
    print(sentences)
    tokenized_sentences = [list(jieba.cut(sentence)) for sentence in sentences]
    vocabulary = set()
    for sentence in tokenized_sentences:
        vocabulary.update(sentence)
    vocabulary = list(vocabulary)
    word_to_index = {word: i for i, word in enumerate(vocabulary)}
    bag_of_words = np.zeros((len(sentences), len(vocabulary)))
    for i, sentence in enumerate(tokenized_sentences):
        for word in sentence:
            word_index = word_to_index[word]
            bag_of_words[i, word_index] += 1

    # 计算词频分布
    word_frequency = bag_of_words / np.linalg.norm(bag_of_words, axis=1, keepdims=True)

    # 保留第一个句子
    filtered_sentences = [sentences[0]+punctuation_list[0]]  # 保留第一个句子
    for i in range(1, len(sentences)):
        similarity = cosine_similarity(word_frequency[i-1], word_frequency[i])
        print(similarity)
        if similarity < 0.75:
            filtered_sentences.append(sentences[i]+punctuation_list[i])
            print(punctuation_list[i])

    # 输出结果

    sentence = ''.join(filtered_sentences)
    return sentence
if __name__ == '__main__':

    print(do_sentence(""))
my_utils.py
tianclll/Ace-Translate