利用虾检测模型检测图像生成的txt转成roLabelImg需要的xml

前沿

前些日子，做了一个活虾检测的模型,一开始是用
roLabelImg进行标注的，奈何虾的个数太多，只用了部分数据。

如果是更小的虾苗，更密集，因此标注是一件很麻烦的事情。在使用了部分数据进行标注并训练之后，得到了用检测虾的模型，这个时候就可以利用训练好的模型在未进行标注的样本进行检查并进行批量标注了。

TXT到xml数据格式的转换

在用虾苗检测模型下进行批量样本检测的时候，生成的目标结果是四个点的坐标，但是不能确保漏检跟误检，所以还需要roLabelImg进行确认。因此需要转成roLabelImg对应的xml文件，roLabelImg对应的文件格式是中心点坐标和宽高，以及旋转角度。因此需要四个点转中心点+旋转角度的转换。
用检测器检测生成的txt格式如下：

每一行对应的是四点坐标x1,y1,x2,y2,x3,y3,x4,y4,class,difficulty

# *_* coding : UTF-8 *_*
# 功能描述   ：把四点坐标x1,y1,x2,y2,x3,y3,x4,y4,class,difficulty转换成旋转框 cx,cy,w,h,angle换成
 
import os
import xml.etree.ElementTree as ET
import math
import cv2
import numpy as np 
'''xml文件的格式'''
head ='''<annotation verified="no">
    <folder>%(folder)s</folder>
    <filename>%(name)s</filename>
    <path>%(path)s</path>
    <source>
        <database>Unknown</database>
    </source>
    <size>
        <width>%(width)d</width>
        <height>%(height)d</height>
        <depth>%(depth)d</depth>
    </size>
    <segmented>0</segmented>
'''
source = '''    <object>
        <type>robndbox</type>
        <name>%(class)s</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <robndbox>
            <cx>%(cx)0.4f</cx>
            <cy>%(cy)0.4f</cy>
            <w>%(w)0.4f</w>
            <h>%(h)0.4f</h>
            <angle>%(angle)0.6f</angle>
        </robndbox>
    </object>
'''
 
end = '''</annotation>
'''
def get_dataset_dirs(input_dir):
  person_names = []
  _subdirs = []
  files = []
  for person_name in os.listdir(input_dir):
    person_names.append(person_name)
  for person_name in person_names:
    _subdir = os.path.join(input_dir, person_name)
    if not os.path.isdir(_subdir):
      continue
    _subdirs.append(_subdir)
  return _subdirs
def txt2xml(txt_path,image_path,xml_save_path):
    img_list_dir = get_dataset_dirs(image_path)
    image_path_list = []
    for dir in img_list_dir:
        for img in os.listdir(dir):
            if not img.endswith('.jpg') and not img.endswith('.png') and not img.endswith('.bmp') and not img.endswith('.jpeg'):
                print (img)
                continue
            image_path_list.append(os.path.join(dir, img))
    source_info = {}
    label_info = {}
    for sub_img_path in image_path_list:    
        folder,image_name = os.path.split(sub_img_path)
        sub_txt_path = txt_path + (image_name.replace('.jpg','.txt'))
        sub_xml_path = xml_save_path + (image_name.replace('.jpg','.xml'))  
        image = cv2.imread(sub_img_path)
        image_w ,image_h,image_c = image.shape

        w_xml = open(sub_xml_path,'w')
        filename = image_name.split('.')[0]
        folder = os.path.basename(folder)
        source_info['name'] = filename
        source_info['folder'] = folder
        source_info['path'] = sub_img_path

        source_info['width'] = image_w
        source_info['height'] = image_h
        source_info['depth'] = image_c
        w_xml.write(head % source_info)
        f = open(sub_txt_path,'r')
        txt_lines = f.readlines()
        for txt_line in txt_lines:           
            tl = txt_line.strip().split(' ')
            cnt = np.array(np.float32([[float(tl[0]),float(tl[1])],[float(tl[2]),float(tl[3])],[float(tl[4]),float(tl[5])],[float(tl[6]),float(tl[7])]]))
            rect = cv2.minAreaRect(cnt)
            theta = rect[-1] #使用 cv2.minAreaRect 计算的角度不准
            box_w = rect[1][0]
            box_h = rect[1][1]
            center_x = rect[0][0]
            center_y = rect[0][1]
  
            tem_x = float(tl[2])-float(tl[0])
            tem_y = float(tl[3])-float(tl[1])
            angle = math.atan(math.fabs(tem_y/tem_x))#额外重新计算角度
            if tem_x<0 and tem_y >=0:#
                angle = math.pi - angle
                box_w,box_h = box_h,box_w#根据不同的位置进行角度转换，以及长短边的转换
            elif(tem_x<0 and tem_y <0):
                angle = math.pi + angle
            elif(tem_x>=0 and tem_y <0):
                angle = math.pi*2.0 - angle
            label_info['cx'] = center_x
            label_info['cy'] = center_y
            label_info['w'] = box_w
            label_info['h'] = box_h
            label_info['angle'] = angle
            label_info['class'] = str(tl[8])
            w_xml.write(source % label_info)
        w_xml.write(end)
        print('write {} done!'.format(sub_xml_path))
if __name__ == '__main__':
    txt_path = "./txt/"#生成的txt路径目录
    image_path = "./Shrimp_/"#对应的图像样本目录
    xml_save_path = "./save_xml/"#要保存的xml文件路径
    txt2xml(txt_path,image_path,xml_save_path)

代码对应的资源包：txt转成roLabelImg需要的xml
在这里，使用opencv的cv2.minAreaRect()函数来计算中心点坐标以及宽高，旋转角度。但是值得注意的是，用这个方法得到的角度跟之前手动标注的是有差异的，需要根据宽高做一些转换（但是由于本人在标注的时候可能有些目标旋转的方向不正确，导致了小部分的角度不准确，跟手动标注的对不上），因此本人没有利用cv2.minAreaRect()生成的角度，而是额外做了角度的换算，见代码注释。

最后

猜想最后的角度有部分不准可能是在手动标注的时候旋转的方向不一样造成的，所以在后续的roLabelImg确认过程中，需要进行调整。附上虾检测的结果图。