在目标检测寻找数据集的过程中可能发现有的图像分割数据集可以满足自己的要求,却没有相应的目标检测格式的标签,接下来是转换代码:
import cv2
import numpy as np
import pandas as pd
import os
def get_coor(img):
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # 变为灰度图
# cv2.imshow('1', gray)
# cv2.waitKey(0)
ret, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU) ## 阈值分割得到二值化图片
# cv2.imshow('2', binary)
contours, heriachy = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
print(len(contours))
arr = np.empty(shape=[0,4])
print(arr)
for i, contour in enumerate(contours):
#if i == len(contours) - 1:
#print('i:', i)
a = sorted(contour[:, 0], key=lambda x: x[0]) # 所有坐标按x轴从小到大排序
x_min = a[0][0]
x_max = a[-1][0]
b = sorted(contour[:, 0], key=lambda x: x[1]) # 所有坐标按y轴从小到大排序
y_min = b[0][1]
y_max = b[-1][1]
rec = img
# cv2.drawContours(img, contours, i, (0, 0, 255), 5)
# 第一个参数指在哪幅图上绘制轮廓信息,第二个参数是轮廓本身,第三个参数是指定绘制哪条轮廓
# 第四个参数是绘图的颜色,第五个参数是绘制的线宽 输入-1则表示填充
cv2.rectangle(rec, (x_min, y_min), (x_max, y_max), (0, 255, 0), 3)
# cv2.imshow('rectangle', rec)
# cv2.waitKey()
a = np.array([str(x_min), str(y_min), str(x_max), str(y_max)])
arr = np.append(arr, [a], axis=0)
return arr
def save_xml(src_xml_dir, img_name, h, w, arr):
xml_file = open((src_xml_dir + '/' + img_name + '.xml'), 'w')
xml_file.write("<annotation>\r\n")
xml_file.write(" <folder>VOC2007</folder>\r\n")
xml_file.write('" <filename>' + str(img_name) + '.jpg' + '</filename>\r\n')
xml_file.write(" <size>\r\n")
xml_file.write(' <width>' + str(w) + '</width>\r\n')
xml_file.write(' <height>' + str(h) + '</height>\r\n')
xml_file.write(' <depth>3</depth>\r\n')
xml_file.write(' </size>\r\n')
# d = str(arr[0][0])
# f = str(arr[1][1])
for i in (range(len(arr))):
# write the region of image on xml file
print("i:", i)
xml_file.write(' <object>\r\n')
xml_file.write(' <name>' + 'strawberry' + '</name>\r\n')#你自己的类别名
xml_file.write(' <pose>Unspecified</pose>\r\n')
xml_file.write(' <truncated>0</truncated>\r\n')
xml_file.write(' <difficult>0</difficult>\r\n')
xml_file.write(' <bndbox>\r\n')
xml_file.write(' <xmin>' + str(arr[i][0]) + '</xmin>\r\n')
xml_file.write(' <ymin>' + str(arr[i][1]) + '</ymin>\r\n')
xml_file.write(' <xmax>' + str(arr[i][2]) + '</xmax>\r\n')
xml_file.write(' <ymax>' + str(arr[i][3]) + '</ymax>\r\n')
xml_file.write(' </bndbox>\r\n')
xml_file.write(' </object>\r\n')
xml_file.write('</annotation>')
file_dir = r'H:\StrawDI_Db1\train\label'#图像分割的png格式标签
save_xml_dir = r'H:\StrawDI_Db1\train\voc'#输出的voc格式的标签
for name in os.listdir(file_dir):
print(name)
img_path = os.path.join(file_dir, name)
img = cv2.imread(img_path)
# img = cv2.imdecode(np.fromfile(img_path), dtype=np.uint8), -1)
h, w = img.shape[:-1]
#x1, y1, x2, y2 = get_coor(img)
arr1 = get_coor(img)
print(arr1)
print(arr1[0][3])
img_name = name.split('.')[0]
save_xml(save_xml_dir, img_name, h, w, arr1)
注意,这是比较粗糙的转换,针对一幅图片里面多个同类别的转换,还没实现多个类别的转换,有时候多个物体重叠会被判定为一个框,需要自己后期检查更正。
参考文章:(21条消息) 分割数据集label转换为目标检测boundingbox_分割标签转目标检测格式_TryRestart的博客-CSDN博客