使用 golang xml 标准库转换 SeaShips(7000) 数据集 标注文件格式 到 YOLO数据格式

使用 golang xml 标准库 160 行代码转换 7000 张图片的标注文件。

rect => <object-class> <x_center> <y_center> <width> <height>

最近需要使用 yolov4 训练该数据集，然后输出一个船舶识别数据模型，所以写了一个小小工具来进行标注文件格式转换。转换后的效果可以使用上次写的标注工具查看一下。

package main

import (
	"encoding/xml"
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
)

// 定义一些 xml 结构和 SeaShips(7000) 数据集对应
type Size struct {
	Width int `xml:"width"`
	Height int `xml:"height"`
	Depth int `xml:"depth"`
}

type BndBox struct {
	Xmin int `xml:"xmin"`
	Ymin int `xml:"ymin"`
	Xmax int `xml:"xmax"`
	Ymax int `xml:"ymax"`
}

type Object struct {
	Name string `xml:"name"`
	Pose string `xml:"pose"`
	Truncated string `xml:"truncated"`
	BndBox BndBox `xml:"bndbox"`
}

type Annotation struct {
	XMLName xml.Name `xml:"annotation"`
	Folder string `xml:"folder"`
	Filename string `xml:"filename"`
	Path string `xml:"path"`
	Size Size `xml:"size"`
	Segmented int `xml:"segmented"`
	Object []Object `xml:"object"`
}

// 获取 SeaShips(7000) 数据集标签对应的下标索引
func GetClassesIndex(classes string) int {
	pair := []string{"bulk cargo carrier","container ship","fishing boat","general cargo ship","ore carrier","passenger ship"}
	switch classes {
	case pair[0]:
		return 0
	case pair[1]:
		return 1
	case pair[2]:
		return 2
	case pair[3]:
		return 3
	case pair[4]:
		return 4
	case pair[5]:
		return 5
	default:
		break
	}
	log.Fatal("not classes ",classes)
	return 0
}

// 计算中心点 boxes
// <object-class> <x_center> <y_center> <width> <height>
func calcCenterBoxes(filename string,in Annotation){
	// .xml 后缀格式文件名 => .txt 后缀格式文件名
	fs,err := os.Create(strings.Split(filename,".")[0] + ".txt")
	defer fs.Close()
	if err != nil {
		log.Fatal(err)
	}

	// 枚举 xml 文件中 boxes 对象信息
	for _,object := range in.Object {
		id := GetClassesIndex(object.Name)
		xCenter := float64(object.BndBox.Xmax + object.BndBox.Xmin) / 2.0 / float64(in.Size.Width)
		yCenter := float64(object.BndBox.Ymax + object.BndBox.Ymin) / 2.0 / float64(in.Size.Height)
		width := float64(object.BndBox.Xmax - object.BndBox.Xmin) / float64(in.Size.Width)
		height := float64(object.BndBox.Ymax - object.BndBox.Ymin) / float64(in.Size.Height)
		// 写入一行数据
		fs.WriteString(fmt.Sprintf("%d %.6f %.6f %.6f %.6f\n",id,xCenter,yCenter,width,height))
	}
}

type Option struct {
	InputPath string
	OutputPath string
	OutPutClasses bool
}

var global Option

func init() {
	flag.StringVar(&global.InputPath, "input .xml path","", "input .xml file path")
	flag.StringVar(&global.OutputPath, "output .txt path","", "input .txt file path")
	flag.BoolVar(&global.OutPutClasses,"print classes",false,"")
}

func main() {
	// 解析命令行参数
	flag.Parse()
	if !flag.Parsed() {
		println("please input *.exe -h get help.")
		os.Exit(0)
	}

	// 创建路径所需要的所有目录
	err := os.MkdirAll(global.OutputPath,os.ModeDir)
	if err != nil {
		log.Fatal(err)
	}

	// 读取目录下的所有文件
	files, err := os.ReadDir(global.InputPath)
	if err != nil {
		log.Fatal(err)
	}

	// 模拟 set 结构存储 xml 中的标签
	set := make(map[string]struct{})
	var exists = struct{}{}

	// 遍历所有文件
	for _, file := range files {
		// 打印文件名称
		fmt.Println(file.Name())
		// 打开 xml 文件
		fs,err := os.Open(global.InputPath + "/" + file.Name())
		if err != nil {
			log.Fatal(err)
		}

		decode := xml.NewDecoder(fs)
		a := Annotation{}
		// 解析 xml 文件
		err = decode.Decode(&a)
		if err != nil {
			log.Fatal(err)
		}

		// 计算 boxes
		calcCenterBoxes(global.OutputPath + "/" + file.Name(),a)

		// 存储标签
		for _, object := range a.Object {
			set[object.Name] = exists
		}
	}

	// 打印标签
	if global.OutPutClasses {
		fmt.Println()
		for k := range set {
			fmt.Println(k)
		}
	}
}