目录
一、Surya模型检测使用Python接口中的源码详解
使用surya源码进行模型检测的过程中， 模型的各种参数设置、环境变量配置都写在下载的源码下 /surya/settings.py 文件中，修改其中的参数即可实现全局的配置。
1.选择模型检测GPU
修改 settings.py 文件中的 TORCH_DEVICE 参数，默认为 None 时，运行监测代码会自动检查当前设备，并选择索引顺序最小的GPU运行——‘cuda：0’。在实际部署中，如果服务器中第一块GPU——‘cuda：0’有其他模型在跑，可能需要调整模型预测位置，将模型放到另一块GPU上运行。只需修改以下参数代码。
# 指定模型所在GPU
TORCH_DEVICE: Optional[str] = 'cuda:1'
2.配置加载模型参数
在 settings.py 文件中模型参数默认为在线加载，当服务器无法连接外部网络时，离线部署加载模型参数需调整设置中的地址。需根据自己模型存放位置，修改为下面的参数字符串。
## 地址需改为你存放模型的绝对地址
# 文本行检测模型
DETECTOR_MODEL_CHECKPOINT: str = "//Surya-OCR/hugging_model/surya_det2"
# 文本区域检测模型
LAYOUT_MODEL_CHECKPOINT: str = "//Surya-OCR/hugging_model/surya_layout"默认在线加载模型参数位置：
参数修改内容：（为你存放离线下载模型地址）
测试模型是否加载成功的代码如下：
from surya.model.detection.segformer import load_model, load_processor
from surya.settings import settings
# 行检测模型：surya_det_2
det_model = load_model()
det_processor = load_processor()
print('det2_model load success')
# 区域检测：surya_layout
model = load_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
processor = load_processor(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
print('layout_model load success')修改后，检验模型加载情况：
3.批量检测图片
实际部署中，官方文档提供 Python 接口检测的代码对单个图片检测顺利，但在批量检测图片集——文件夹时报错。 （单图代码写在上一篇博文《从零开始使用Surya-OCR——文本目标检测模型的安装与部署》中，具体可参考https://blog.csdn.net/qq_58718853/article/details/137150986）
第一个报错是，官方提供接口函数，无法读取文件夹内图片，报读取文件权限被拒。暂未实现直接解决该问题的办法。参看后续 batch_text_detection 源代码传参信息，得知图片读取后的传入函数结果是一个列表，可以选择替代方案实现同等效果，替代方案代码如下。
import os
from PIL import Image
from surya.detection import batch_text_detection
from surya.model.detection.segformer import load_model, load_processor
IMAGE_PATH = 'image_path'
model = load_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
processor = load_processor(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
print('model load success')
# 批量将文件夹图片读入images列表中
images = []
for file in os.listdir(IMAGE_PATH):
image_path = os.path.join(IMAGE_PATH, file)
image = Image.open(image_path)
images.append(image)
predictions = batch_text_detection(images, model, processor)
print(predictions)
使用新代码后，原本的报错问题解决了，但出现了新的报错。
第二个报错是，检查surya模型做批量预测任务时，在得到模型输出后，还会对多张图片的结果进行多进程处理。问题就在多进程处理时，源代码内置函数重复调用主函数——导致模型本来只需加载一次，此时不断加载模型致使GPU显存爆了。我们找到对应报错的源码处检查。
按住 Ctrl 并点击报错的 batch_text_detection 函数，即可进入源码处，检查研究发现可能是Windows系统和 Linux 系统对于多进程和多线程的解释存在差异，本机是windows系统， 可能在导入 ProcessPoolExecutor 函数时，将我的主函数视为多进程对象，不断创建新的进程。但是此处希望实现的是中间过程结果的多线程处理，而不影响主函数。因此需要将 ProcessPoolExecutor 改为使用多线程的 ThreadPoolExecutor，问题即可解决。 
经过上述所有源代码修改处理，成功运行主函数，得到surya模型批量检测图片后得到的框信息结果，并将其打印出来。
4.检测输出结果源码解读
Surya模型输出是自定义类的数据格式，下面根据其官方文档和项目源码解读其输出的格式，以方便后续对输出的处理，提取出所需的数据信息。
官方文档：https://github.com/VikParuchuri/surya
Surya 模型有三种预测模式——OCR & Text Line & Layout，对应三种模型输出的格式，每种模式的输出都是以类的形式定义的。下面重点放在 Text Line 文本行检测和 Layout 区域检测的源码信息解读上。
①文本行检测的模型输出——Text Line
将与输出相关的源代码从项目中单独提取出来看，下面是输出的基础类，即每个图片模型预测后的信息都封装在了 TextDetectionResult 里面。
"文本行检测"
# 输出的基础类
class TextDetectionResult(BaseModel):
bboxes: List[PolygonBox]
vertical_lines: List[ColumnLine]
horizontal_lines: List[ColumnLine]
heatmap: Any
affinity_map: Any
image_bbox: List[float]
下面分别解释输出基础类中的具体信息都是怎么定义的，从源码中提出相关代码。
输出的第一个类信息：PolygonBox 注解
（下面非完整代码，为清晰类输出含义，只选取主要功能代码）
# 输出框信息类
class PolygonBox(BaseModel):
polygon: List[List[float]] ## 存储框四个角——全坐标
confidence: Optional[float] = None ## 框预测置信度
def bbox(self) -> List[float]:
box = [self.polygon[0][0], self.polygon[0][1], self.polygon[1][0], self.polygon[2][1]]
if box[0] > box[2]:
box[0], box[2] = box[2], box[0]
if box[1] > box[3]:
box[1], box[3] = box[3], box[1]
return box ## 存储框左上右下——对角坐标通过源码可知，此类保存的是检测框的坐标信息和置信度，这是预测中的主要信息。使用具体的模型输出结果，可以更清楚该类的输出形状。输出TextDetectionResult 包含多个类信息，其中定义的 bboxes ——即 PolygonBox 存储的是一张图片内检测出来的所有框，而每个框的信息结构是包含三个子类：全坐标、置信度和对角坐标。
输出的第二个类信息：ColumnLine 注解
（下面非完整代码，为清晰类输出含义，只选取主要功能代码）
# 图片中线的检测
class ColumnLine(Bbox):
vertical: bool # 垂直线：有-True;无-False
horizontal: bool # 水平线：有-True;无-False
# 检测线的框坐标
class Bbox(BaseModel):
bbox: List[float]
同样通过源码可知，输出的第二个类信息 ColumnLine 是用来保存在模型预测中对图片检测到的水平线、垂直线的信息，如果检测到了那么就同时保存其框位置。通过具体模型输出对应位置检索，可以清晰理解。
输出的剩余类信息：heatmap、affinity_map、image_bbox
剩余的类信息是对图片结果输出的补充，我们可以直接将其打印出来，看看其内容。最容易看出来的是 image_bbox ，其实际就是用一个框将整个图片框起来，然后返回对角坐标。而heatmap、affinity_map 则是 PIL.Image.Image 的类信息。
②文本区域检测的模型输出 ——Layout
同文本行检测一样，废话少说，直接上代码和图。
# 区域检测输出
class LayoutResult(BaseModel):
bboxes: List[LayoutBox]
segmentation_map: Any
image_bbox: List[float]、
class LayoutBox(PolygonBox):
label: str ## 多了一个区域类别的预测
5.批量信息的保存和可视化
将surya模型自定义的类输出转化为 json 格式保存到指定文件夹，代码如下。
import os
import json
from PIL import Image
from surya.detection import batch_text_detection
from surya.model.detection.segformer import load_model, load_processor
import cv2
import numpy as np
IMAGE_PATH = 'iamge_path' ## 检测图片保存地址
json_file = 'json_path' ## 框json保存地址
checkpoint = 'model_path' ## 模型参数加载地址
heat_file = 'heat_path' ## 热图保存
########### 上述为修改部分，根据实际地址填入，下面无需修改 ############
model, processor = load_model(checkpoint=checkpoint), load_processor(checkpoint=checkpoint)
print('model load success')
# 模型预测
images = []
image_name = []
for file in os.listdir(IMAGE_PATH):
image_path = os.path.join(IMAGE_PATH, file)
image = Image.open(image_path)
images.append(image)
image_name.append(file)
predictions = batch_text_detection(images, model, processor)
print('predict success')
# 保存模型结果
## 类型转为json
def class_to_json(bboxes, file, box_type=True):
json_list = []
for i, bbox in enumerate(bboxes):
if box_type:
json_dict = dict()
box = bbox.bbox
box.append(bbox.confidence)
json_dict["id"] = i
json_dict["name"] = file
json_dict["box"] = box
json_list.append(json_dict)
else:
json_dict = dict()
box = bbox.bbox
json_dict["id"] = i
json_dict["name"] = file
json_dict["box"] = box
json_list.append(json_dict)
return json_list
## 保存到指定文件夹
def save_json(json_list, json_path):
with open(json_path, 'w') as f:
json.dump(json_list, f)
## 主函数
def save_predict(predictions, image_name, heat_file):
for i, pred in enumerate(predictions):
# 框信息保存
bboxes = pred.bboxes
vertical = pred.vertical_lines
horizontal = pred.horizontal_lines
file = image_name[i]
bboxes_json = class_to_json(bboxes, file)
vertical_json = class_to_json(vertical, file, box_type=False)
horizontal_json = class_to_json(horizontal, file, box_type=False)
basename = file.split('.')[0]
save_json(bboxes_json, os.path.join(json_file+'box/', basename + '.json'))
save_json(vertical_json, os.path.join(json_file + 'vertical/', basename + '.json'))
save_json(horizontal_json, os.path.join(json_file + 'horizontal/', basename + '.json'))
# 热图调参信息保存
heatmap = pred.heatmap
img = cv2.cvtColor(np.asarray(heatmap), cv2.COLOR_RGB2BGR)
cv2.imwrite(heat_file+basename+'.jpg', img)
print(basename + ' success')
if __name__ == '__main__':
save_predict(predictions, image_name, heat_file)可视化框的代码如下。
import os
import json
import cv2
# jpg、json、vis文件位置
jpg_path = 'JPG'
json_path = 'JSON'
vis_path = 'VIS'
########### 上述为修改部分，根据实际地址填入，下面无需修改 ############
# 可视化锚框
## 锚框展示细节
def hsv2bgr(h, s, v):
h_i = int(h * 6)
f = h * 6 - h_i
p = v * (1 - s)
q = v * (1 - f * s)
t = v * (1 - (1 - f) * s)
r, g, b = 0, 0, 0
if h_i == 0:
r, g, b = v, t, p
elif h_i == 1:
r, g, b = q, v, p
elif h_i == 2:
r, g, b = p, v, t
elif h_i == 3:
r, g, b = p, q, v
elif h_i == 4:
r, g, b = t, p, v
elif h_i == 5:
r, g, b = v, p, q
return int(b * 255), int(g * 255), int(r * 255)
def random_color(id):
h_plane = (((id << 2) ^ 0x937151) % 100) / 100.0
s_plane = (((id << 3) ^ 0x315793) % 100) / 100.0
return hsv2bgr(h_plane, s_plane, 1)
# 可视化主函数
def visualize(json_path, jpg_path, vis_path, box_type=True):
if box_type:
for file in os.listdir(json_path):
with open(json_path+file,'r') as f:
drawResult = json.load(f)
basefile = file.split('.')[0]
jpg_file = os.path.join(jpg_path,basefile+".jpg")
img = cv2.imread(jpg_file)
for idx, result in enumerate(drawResult):
left, top, right, bottom = int(result['box'][0]), int(result['box'][1]), int(result['box'][2]), int(result['box'][3])
label = int(result['box'][4])
color = random_color(1)
cv2.rectangle(img, (left, top), (right, bottom), color=color ,thickness=2, lineType=cv2.LINE_AA)
caption = f"{'ZW'}"
w, h = cv2.getTextSize(caption, 0, 1, 2)[0]
cv2.rectangle(img, (left - 3, top - 33), (left + w + 10, top), color, -1)
cv2.putText(img, caption, (left, top - 5), 0, 1, (0, 0, 0), 2, 16)
save_file = os.path.join(vis_path, basefile+".jpg")
print(save_file)
cv2.imwrite(save_file, img)
else:
for file in os.listdir(json_path):
with open(json_path+file,'r') as f:
drawResult = json.load(f)
basefile = file.split('.')[0]
jpg_file = os.path.join(jpg_path,basefile+".jpg")
img = cv2.imread(jpg_file)
for idx, result in enumerate(drawResult):
left, top, right, bottom = int(result['box'][0]), int(result['box'][1]), int(result['box'][2]), int(result['box'][3])
color = random_color(1)
cv2.rectangle(img, (left, top), (right, bottom), color=color ,thickness=2, lineType=cv2.LINE_AA)
caption = f"{'ZW'}"
w, h = cv2.getTextSize(caption, 0, 1, 2)[0]
cv2.rectangle(img, (left - 3, top - 33), (left + w + 10, top), color, -1)
cv2.putText(img, caption, (left, top - 5), 0, 1, (0, 0, 0), 2, 16)
save_file = os.path.join(vis_path, basefile+".jpg")
print(save_file)
cv2.imwrite(save_file, img)
if __name__ == '__main__':
visualize(json_path+'box/', jpg_path, vis_path+'box/')
visualize(json_path + 'vertical/', jpg_path, vis_path + 'vertical/', box_type=False)
visualize(json_path + 'horizontal/', jpg_path, vis_path + 'horizontal/', box_type=False)
942
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
