[二] pytorch + PyQt5 制作一个图像检测及保存数据集功能
pytorch+pyqt5 制作一个基于yolov5的图像检测功能,并可将识别结果转为数据集
·
经过迭代,该版本已经作废,新版本请看 [一] 基于PyQt的图像算法开发软件——如何让你开发效率翻倍
目的
在完成demo的基本功能后,开始想着拓展其他功能,考虑到之前有调试yolov5(pytorch版本)的经验,如果能用UI更加直接的展示其结果,那调试起来也会相对舒服一点。此外,如果可以把得到了检测结果再保存成数据集,那就可以大大减少人力。
需要完成的目标:
一、可以加载模型文件(.pt),并检测图像内容。
二、可以调整后处理参数(置信区间和交并比阈值)
三、对于多分类模型,可以自定义选择想要识别出的物体类别,并将其再次存储为数据集(.xml)
一、使用效果
后处理演示
视频里展示了该功能的演示
流程为: 导入模型和打开图片->选择想要识别的类->点击识别->调整置信区间和交并比阈值->保存数据
二、部分说明
这里重点介绍下保存数据的过程, 如果有调试过yolov5,那应该就会知道数据集的预处理是先得把图像和对应的xml进行转化,同理,我们也可以把检测的结果再转回图像和对应的xml文件.
部分代码说明:
def create_xml(xml_template, save_path, img, bbox, abspath=True):
if img is not None:
if abspath:
save_path = os.path.abspath(save_path)
data_path = os.path.join(save_path, 'data')
annotations_path = os.path.join(save_path, 'Annotations')
if not os.path.exists(data_path):
os.mkdir(data_path)
if not os.path.exists(annotations_path):
os.mkdir(annotations_path)
img_name = len(os.listdir(data_path))
dom = xml.dom.minidom.parse(xml_template)
folder = dom.getElementsByTagName("folder")[0]
folder.childNodes[0].data = save_path
filename = dom.getElementsByTagName("filename")[0]
filename.childNodes[0].data = '{}.jpg'.format(img_name)
path = dom.getElementsByTagName("path")[0]
path.childNodes[0].data = os.path.join(data_path, '{}.jpg'.format(img_name))
database = dom.getElementsByTagName("database")[0]
database.childNodes[0].data = 'Unknown'
width = dom.getElementsByTagName("width")[0]
height = dom.getElementsByTagName("height")[0]
depth = dom.getElementsByTagName("depth")[0]
if len(img.shape) < 3:
width.childNodes[0].data, height.childNodes[0].data, depth.childNodes[0].data = img.shape[0], img.shape[
1], 1
else:
width.childNodes[0].data, height.childNodes[0].data, depth.childNodes[0].data = img.shape
segmented = dom.getElementsByTagName("segmented")[0]
segmented.childNodes[0].data = 0
root = dom.getElementsByTagName("annotation")[0]
for i in range(len(bbox)):
o = dom.createElement("object")
name = dom.createElement("name")
name.appendChild(dom.createTextNode(str(bbox[i].get('name', 'Unknown'))))
o.appendChild(name)
pose = dom.createElement("pose")
pose.appendChild(dom.createTextNode(str(bbox[i].get('pose', 'Unspecified'))))
o.appendChild(pose)
truncated = dom.createElement("truncated")
truncated.appendChild(dom.createTextNode(str(bbox[i].get('truncated', '0'))))
o.appendChild(truncated)
difficult = dom.createElement("difficult")
difficult.appendChild(dom.createTextNode(str(bbox[i].get('difficult', '0'))))
o.appendChild(difficult)
bndbox = dom.createElement("bndbox")
xmin = dom.createElement("xmin")
xmax = dom.createElement("xmax")
ymin = dom.createElement("ymin")
ymax = dom.createElement("ymax")
xmin.appendChild(dom.createTextNode(str(bbox[i].get('xmin', '0'))))
xmax.appendChild(dom.createTextNode(str(bbox[i].get('xmax', '0'))))
ymin.appendChild(dom.createTextNode(str(bbox[i].get('ymin', '0'))))
ymax.appendChild(dom.createTextNode(str(bbox[i].get('ymax', '0'))))
bndbox.appendChild(xmin)
bndbox.appendChild(ymin)
bndbox.appendChild(xmax)
bndbox.appendChild(ymax)
o.appendChild(bndbox)
root.appendChild(o)
cv2.imwrite(os.path.join(data_path, '{}.jpg'.format(img_name)), img)
with open(os.path.join(annotations_path, "{}.xml".format(img_name)), 'w') as fh:
dom.writexml(fh)
if __name__ == '__main__':
img = cv2.imread('15.jpg')
bbox = [{'name': 'halobolite', 'xmin': 8, 'xmax': 493, 'ymin': 161, 'ymax': 452},
{'name': 'halobolite', 'xmin': 507, 'xmax': 969, 'ymin': 175, 'ymax': 528}]
create_xml('./example.xml', '.', img, bbox) # 需要一个模板文件

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。
更多推荐
所有评论(0)