在给模型传入数据时出现问题(经过调试,数据已经传入了模型处理,但是到了某一步的时候变成了None):
代码如下:
import torch
import json
import base64
from seal_bak.seal_recognition import work
import os
import numpy as np
import cv2
import matplotlib.pyplot as plt
# root_path = os.getcwd()
# model_path = root_path + '/models/seal_detect_best.pt'
model_path = 'D:/test11/yolov5/models/seal_detect_best.pt'
class Infer_main:
def __init__(self, yolo_model_path):
self.model_path = yolo_model_path
self.seal_infer = self.yolo_model_load()
def yolo_model_load(self):
model = torch.hub.load('D:/test11/yolov5',
'custom',
path='D:/test11/yolov5/models/seal_detect_best.pt',
source='local',
force_reload=True) # local repo
model.conf = 0.4
model.eval()
return model
def predict(self, images):
result = {
"base64": images,
"debug": False # debug模式将可视化各环节,否则只输出结果
}
# for item, i in enumerate(images):
yolo_res = self.seal_infer(images)
print(f'yolo_res:{yolo_res}')
yolo_res_list = yolo_res.pandas().xyxy[0].values.tolist()
if yolo_res_list:
for i in range(len(yolo_res_list)):
# print(yolo_res_list[i])
yolo_res_liststr = json.dumps(yolo_res_list[i])
# print(type(yolo_res_liststr))
base64str = base64.b64encode(yolo_res_liststr.encode('utf-8'))
# img_array = np.fromstring(base64str, np.uint8)
# print(type(base64str))
result['base64'] = base64str
result_rec = work(result)
print(result_rec)
return result_rec
if __name__ == '__main__':
from PIL import Image
#
# model_path = 'D:/test11/yolov5/models/seal_detect_best.pt'
# image = plt.imread('D:/test11/yolov5/8.png')
# image = cv2.imread('D:/test11/yolov5/8.png')
image = Image.open('D:/test11/yolov5/1.png')
a = Infer_main(model_path)
# print()
# a.seal_infer(image).show()
a.predict(image)
# model = yolo_model_load(model_path)
# yolo_res = model(image, size=640)
# yolo_res.show()
# yolo_res_list = yolo_res.pandas().xyxy[0].values.tolist()
# print(type(yolo_res_list))
# print('印章识别结果:', yolo_res_list)
报错信息如下:
Traceback (most recent call last):
File "D:\test11\yolov5\test1.py", line 66, in <module>
a.predict(image)
File "D:\test11\yolov5\test1.py", line 51, in predict
result_rec = work(result)
File "D:\test11\seal_bak\seal_recognition.py", line 116, in work
category, img = model.predict(cfg['base64'])
File "D:\test11\seal_bak\utils\Model.py", line 38, in predict
Tensor, img = self.__img2tensor(ImgPath)
File "D:\test11\seal_bak\utils\Model.py", line 58, in __img2tensor
Img = Transfm(OImg)
File "D:\Anaconda3\envs\cj-env\lib\site-packages\torchvision\transforms\transforms.py", line 95, in __call__
img = t(img)
File "D:\Anaconda3\envs\cj-env\lib\site-packages\torchvision\transforms\transforms.py", line 135, in __call__
return F.to_tensor(pic)
File "D:\Anaconda3\envs\cj-env\lib\site-packages\torchvision\transforms\functional.py", line 137, in to_tensor
raise TypeError(f"pic should be PIL Image or ndarray. Got {type(pic)}")
TypeError: pic should be PIL Image or ndarray. Got <class 'NoneType'>
Process finished with exit code 1
该报错为代码逻辑,数据类型不对或没有则返回这段
模型处理代码如下:
class Model():
def __init__(self, choice=2):
if choice == 0:
model = network.SPP_Net1()
Dict = SPPDict1
elif choice == 1:
model = network.SPP_Net2()
Dict = SPPDict2
else:
model = network.SPP_Net3()
Dict = SPPDict3
if self.Device() == "cuda":
model.load_state_dict(torch.load(Dict))
else:
model.load_state_dict(torch.load(Dict, map_location=torch.device('cpu')))
# print("Model loaded!")
self.model = model
def Device(self):
device = "cuda" if torch.cuda.is_available() else "cpu"
return device
def predict(self, ImgPath):
Tensor, img = self.__img2tensor(ImgPath)
with torch.no_grad():
pred = self.model(Tensor)
result = self.__arg2label(pred.argmax())
# print(result)
return result, img
# def __img2tensor(self,ImgPath):
# OImg = cv2.imread(ImgPath) #origin图片
# Transfm = Compose([transforms.ToTensor()])
# Img = Transfm(OImg)
# Img = Img.unsqueeze(0)
# return Img,OImg
def __img2tensor(self, base64_data):
img_b64decode = base64.b64decode(base64_data) # base64解码
# 将base64转换numpy array
img_array = np.fromstring(img_b64decode, np.uint8)
# 转换cv2
OImg = cv2.imdecode(img_array, cv2.COLOR_BGR2RGB)
Transfm = Compose([transforms.ToTensor()])
Img = Transfm(OImg)
Img = Img.unsqueeze(0)
return Img, OImg
def __arg2label(self, arg):
return LabelDict[str(arg.item())]
根据debug调试,是到了__img2tensor函数中出现的问题,在OImg = cv2.imdecode(img_array, cv2.COLOR_BGR2RGB)这步时 OImg得到的是None
模型逻辑代码和所需数据如下:
import re
import time
import traceback
import base64
from seal_bak.utils.ocr import ellipse_ocr, circle_ocr, rectangle_ocr
from seal_bak.utils.detect import *
from seal_bak.utils.Model import Model
import logging
import copy
import numpy as np
model = Model(0)
CODE_SCORE = 0 # 0.98 0.97
NAME_SCORE = 0 # 0.9
def get_result(res, category):
en_or_num = re.compile(r"[a-zA-Z0-9]", re.I) # [a-z]|\d
not_en_nor_num = re.compile(r"[^a-zA-Z0-9]", re.I) # 非英文数字
null_text = {"text": "", "confidence": 0}
result = {
"code": null_text, # 印章编号
"name": null_text, # 印章名称
"SealType": "99", # 印章类型代码
"strSealType": "其他类型印章" # 印章类型名称
}
seal_type = {
"法定名称章": "01",
"财务专用章": "02",
# "发票专用章": "03",
"合同专用章": "04",
"未找到印章": "06",
# "法定代表人名章": "05",
"其他类型印章": "99"
}
# logging.info(f"<{[str(item['text']) + '|' + str(item['confidence']) for item in res]}>")
if category == "正方形":
result["SealType"] = "05"
result["strSealType"] = "法定代表人名章"
result["code"] = {
"text": res[0]["text"], "confidence": res[0]["confidence"]
} if res[0]["confidence"] >= CODE_SCORE else null_text # 设定编号
result["name"] = {
"text": res[1]["text"], "confidence": res[1]["confidence"]
} if res[1]["confidence"] >= NAME_SCORE else null_text # 设定name
elif category == "椭圆": # todo 调参
result["SealType"] = "03"
result["strSealType"] = "发票专用章"
for item in res:
text = item["text"]
score = item["confidence"]
from_ = item["from"]
if len(en_or_num.findall(text)) > len(not_en_nor_num.findall(text)):
result["code"] = {
"text": text, "confidence": score
} if score >= CODE_SCORE and from_ == "center" else null_text # 设定编号
else:
result["name"] = {
"text": text, "confidence": score
} if score >= NAME_SCORE and from_ == "side" else null_text # 设定name
elif len(res) == 2:
result["SealType"] = "01"
result["strSealType"] = "法定名称章"
for item in res:
text = item["text"]
score = item["confidence"]
from_ = item["from"]
if len(en_or_num.findall(text)) > len(not_en_nor_num.findall(text)):
result["code"] = {
"text": text, "confidence": score
} if score >= CODE_SCORE and from_ == "side" else null_text # 设定编号
else:
result["name"] = {
"text": text, "confidence": score
} if score >= NAME_SCORE and from_ == "side" else null_text # 设定name
elif category == False:
result["SealType"] = "06"
result["strSealType"] = "未找到印章"
else:
set_type = True
for item in res:
text = item["text"]
score = item["confidence"]
from_ = item["from"]
if len(en_or_num.findall(text)) > len(not_en_nor_num.findall(text)):
result["code"] = {
"text": text, "confidence": score
} if score >= CODE_SCORE and from_ == "side" else null_text # 设定编号
else:
if set_type and from_ == "center":
# 设定type
if text not in seal_type.keys():
result["SealType"] = "99"
result["strSealType"] = "其他类型印章" # "其他类型印章"
else:
result["SealType"] = seal_type[text]
result["strSealType"] = text
set_type = False
else:
# 设定name
result["name"] = {
"text": text, "confidence": score
} if score >= NAME_SCORE and from_ == "side" else null_text # 设定name
return result
def work(cfg):
t0 = time.time()
# 分类
# category, img = model.predict(cfg['img_path'])
category, img = model.predict(cfg['base64'])
img = enlarge_img(img)
img_ = np.copy(img)
# 提取红色
t1 = time.time()
img_, img_bw = Lazyfilter(img_, cfg)
# 开运算去噪填充
t2 = time.time()
img_ = erode_dilate(img_, category, cfg=cfg)
# 查找最大轮廓
t3 = time.time()
contours, max_idx = find_max(img_)
# 完成检测框
t4 = time.time()
det = get_area(img, contours, max_idx, category, cfg=cfg) # 不进行分类
# 截取目标区域 img:原图 img_bw:二值
t5 = time.time()
img = rotate_cut(img, det, cfg=cfg)
# ============分类处理目标区域
t6 = time.time()
# if category == False:
# empty_ocr()
try:
if category == '圆形':
res = circle_ocr(img, cfg=cfg) # 使用二值图预测
elif category == "正方形":
res = rectangle_ocr(img, cfg=cfg)
else:
res = ellipse_ocr(img, cfg=cfg)
except Exception:
print("未找到印章")
# 将文本列表处理成标准json格式
try:
result = get_result(res, category)
except Exception:
pass
try:
result_string = json.dumps(result, ensure_ascii=False)
except Exception:
pass
if not cfg["debug"]:
# with open(os.path.join(cfg["to_path"], cfg["file_name"] + ".json"), mode="w", encoding="utf-8") as f:
# json.dump(result, f, ensure_ascii=False)
pass
t7 = time.time()
# detail = f"{cfg['file_name']} done in {round(t7 - t0, 2)}s. " \
# f"class={det['class']}, text={result_string}"
# logging.info(detail)
print(t7 - t1)
try:
return result_string
except Exception:
pass
Debug调试的结果如下(其余步骤正常我就省略了):
你说有问题的这句换成下面的看下
OImg = cv2.imdecode(img_array,flags=cv2.IMREAD_COLOR)
OImg =cv2.cvtColor(OImg,cv2.COLOR_BGR2RGB)
将 cv2.imdecode 与 cv2.cvtColor 分开写,分别进行调用:
def __img2tensor(self, base64_data):
img_b64decode = base64.b64decode(base64_data) # base64解码
# 将base64转换numpy array
img_array = np.fromstring(img_b64decode, np.uint8)
# 转换cv2
OImg = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
OImg = cv2.cvtColor(OImg, cv2.COLOR_BGR2RGB)
Transfm = Compose([transforms.ToTensor()])
Img = Transfm(OImg)
Img = Img.unsqueeze(0)
return Img, OImg
将 cv2.COLOR_BGR2RGB 修改为 cv2.IMREAD_COLOR
OImg = cv2.imdecode(img_array, cv2.COLOR_BGR2RGB)
cv2.COLOR_BGR2RGB 用错了吧,cv2.IMREAD_COLOR用在文件处理, cv2.COLOR_BGR2RGB 是颜色处理
试试
OImg = cv2.imdecode(img_array,cv2.IMREAD_COLOR)
OImg =cv2.cvtColor(OImg,cv2.COLOR_BGR2RGB)
引用chatGPT作答,根据报错信息,发现问题出在seal_recognition.py文件中的model.predict(cfg['base64'])语句上。具体来说,model.predict()方法中需要传入的参数应该是一张 PIL Image 或 numpy.ndarray 类型的图片,但是当前传入的是一个 NoneType 类型的变量。
根据代码,在 Infer_main 类中的 predict() 方法中,result['base64'] 的类型是 bytes。然后在 work() 方法中,将 result['base64'] 作为参数传入 model.predict() 方法中,导致了报错。
因此,需要将 result['base64'] 转换为 PIL Image 或 numpy.ndarray 类型的图片,并且确保转换后的图片不为 None。可以使用 PIL 或 opencv 库中的方法来实现这个过程。例如,使用 PIL 库中的 Image.open() 方法将 bytes 类型的数据转换为 PIL Image 类型,或者使用 opencv 库中的 cv2.imdecode() 方法将 bytes 类型的数据转换为 numpy.ndarray 类型。
以下是一个可能的解决方案:
from PIL import Image
import io
class Infer_main:
def predict(self, images):
result = {
"base64": images,
"debug": False
}
image = Image.open(io.BytesIO(result['base64']))
# 或者使用以下代码
# image = cv2.imdecode(np.frombuffer(result['base64'], np.uint8), cv2.IMREAD_COLOR)
result_rec = work(image)
return result_rec
class Model():
def predict(self, img):
img_tensor = self.__img2tensor(img)
...
def __img2tensor(self, img):
transform = T.Compose([
T.Resize((512, 512)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
img_tensor = transform(img).unsqueeze(0)
return img_tensor, img
其中,io.BytesIO 用于创建一个文件对象,用于将 bytes 类型的数据读入到内存中;cv2.imdecode() 方法用于解码图片,并将其转换为 numpy.ndarray 类型的数据。
可以考虑以下几个可能的原因:
数据传入模型的过程中发生了错误,导致数据在某一步骤被改变为 None。可以检查一下传入数据的代码,看看是否有任何可能导致数据出错的地方。
模型本身存在问题,在某一步骤处理数据时出现了错误。可以检查一下模型的代码,看看是否有任何可能导致数据出错的地方。
数据格式不正确,导致模型无法正确处理数据。可以检查一下数据格式是否符合模型要求。
针对这些可能的问题,可以尝试以下解决思路:
检查传入数据的代码,确保传入的数据是正确的。
检查模型的代码,确保模型可以正确处理传入的数据。
检查数据格式是否正确,并根据需要进行数据格式转换。
具体的解决方案需要根据具体情况进行调整。如果可以提供更多的代码和错误信息,可能能够更好地帮助解决问题。