GOT-OCR-2.0-hf开源模型：多语言OCR技术的全面提升

# 安装必要的依赖
# pip install transformers
# pip install torch
# pip install opencv-python

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import cv2
import numpy as np

# 加载模型和分词器
model_path = "got-it-ai/GOT-OCR-2.0-hf"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True
)

# 加载图像
image_path = "example_image.jpg"
image = cv2.imread(image_path)

# 图像预处理
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 执行OCR
result = model.chat(tokenizer, image, ocr_type="ocr")

# 输出结果
print(f"OCR结果: {result}")

# 图像特征提取示例
import torch
import torch.nn as nn
import torchvision.models as models

# 改进的ResNet50-V1d骨干网络
class ModifiedResNet50V1d(nn.Module):
    def __init__(self):
        super(ModifiedResNet50V1d, self).__init__()
        # 加载预训练的ResNet50-V1d
        resnet = models.resnet50(pretrained=True)
        
        # 改进: 替换第一个卷积层，增加dilation
        self.conv1 = nn.Conv2d(
            3, 64, kernel_size=7, stride=2, padding=3, dilation=1, bias=False
        )
        # 复制预训练权重
        self.conv1.weight.data = resnet.conv1.weight.data
        
        # 其余层保持不变
        self.bn1 = resnet.bn1
        self.relu = resnet.relu
        self.maxpool = resnet.maxpool
        self.layer1 = resnet.layer1
        self.layer2 = resnet.layer2
        self.layer3 = resnet.layer3
        self.layer4 = resnet.layer4
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        
        return x

# 使用示例
feature_extractor = ModifiedResNet50V1d()
# 输入形状: (batch_size, channels, height, width)
input_image = torch.randn(1, 3, 224, 224)
features = feature_extractor(input_image)
print(f"特征图形状: {features.shape}")

# 多语言文本识别示例
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载模型和分词器
model_path = "got-it-ai/GOT-OCR-2.0-hf"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True
)

# 支持的主要语言列表
supported_languages = [
    "中文", "英文", "日文", "韩文", "西班牙文", "法文", "德文", 
    "俄文", "阿拉伯文", "葡萄牙文", "意大利文", "荷兰文", "瑞典文",
    "芬兰文", "挪威文", "丹麦文", "波兰文", "捷克文", "匈牙利文"
]

print("GOT-OCR-2.0-hf支持的主要语言:")
for i, lang in enumerate(supported_languages, 1):
    print(f"{i}. {lang}")

# 多语言识别示例
# 假设我们有不同语言的图像
# image_zh = cv2.imread("chinese_text.jpg")
# image_en = cv2.imread("english_text.jpg")
# image_jp = cv2.imread("japanese_text.jpg")

# result_zh = model.chat(tokenizer, image_zh, ocr_type="ocr")
# result_en = model.chat(tokenizer, image_en, ocr_type="ocr")
# result_jp = model.chat(tokenizer, image_jp, ocr_type="ocr")

# print(f"中文识别结果: {result_zh}")
# print(f"英文识别结果: {result_en}")
# print(f"日文识别结果: {result_jp}")

# 文档数字化示例
import cv2
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载模型
model_path = "got-it-ai/GOT-OCR-2.0-hf"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True
)

# 读取文档图像
image = cv2.imread("document.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 执行OCR
result = model.chat(tokenizer, image, ocr_type="ocr")

# 保存为文本文件
with open("digitized_document.txt", "w", encoding="utf-8") as f:
    f.write(result)

print("文档数字化完成，结果已保存到digitized_document.txt")

# 自动化数据提取示例
import cv2
import pandas as pd
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import re

# 加载模型
model_path = "got-it-ai/GOT-OCR-2.0-hf"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True
)

# 读取表单图像
image = cv2.imread("form.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 执行OCR
ocr_result = model.chat(tokenizer, image, ocr_type="ocr")

# 提取表单数据的正则表达式示例
# 假设表单包含姓名、年龄、邮箱字段
name_pattern = r"姓名[:：]\s*(.+?)\n"
age_pattern = r"年龄[:：]\s*(\d+)"
email_pattern = r"邮箱[:：]\s*([a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,})"

# 提取数据
name_match = re.search(name_pattern, ocr_result)
age_match = re.search(age_pattern, ocr_result)
email_match = re.search(email_pattern, ocr_result)

# 构建数据字典
data = {
    "姓名": name_match.group(1) if name_match else "",
    "年龄": int(age_match.group(1)) if age_match else 0,
    "邮箱": email_match.group(1) if email_match else ""
}

# 保存为CSV文件
df = pd.DataFrame([data])
df.to_csv("extracted_data.csv", index=False, encoding="utf-8")

print("数据提取完成，结果已保存到extracted_data.csv")
print("提取的数据:")
print(df)

# 安装必要的依赖
pip install transformers
pip install torch
pip install opencv-python
pip install pandas  # 用于数据处理
pip install numpy

# 下载模型（可选，自动加载）
# from transformers import AutoModelForCausalLM, AutoTokenizer
# model_path = "got-it-ai/GOT-OCR-2.0-hf"
# tokenizer = AutoTokenizer.from_pretrained(model_path)
# model = AutoModelForCausalLM.from_pretrained(model_path)

# 批量处理图像示例
import os
import cv2
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载模型
model_path = "got-it-ai/GOT-OCR-2.0-hf"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True
)

# 图像文件夹路径
image_folder = "images/"
output_folder = "outputs/"

# 创建输出文件夹
os.makedirs(output_folder, exist_ok=True)

# 遍历文件夹中的所有图像
for filename in os.listdir(image_folder):
    if filename.endswith(".jpg") or filename.endswith(".png"):
        # 构建完整路径
        image_path = os.path.join(image_folder, filename)
        
        # 读取图像
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 执行OCR
        result = model.chat(tokenizer, image, ocr_type="ocr")
        
        # 保存结果
        output_filename = os.path.splitext(filename)[0] + ".txt"
        output_path = os.path.join(output_folder, output_filename)
        
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(result)
        
        print(f"处理完成: {filename} -> {output_filename}")

print("\n批量处理完成！")