物体認識用データセットファイル名変更プログラム

11515 ワード
Keras Python3 YOLOv3 TensorFlow ssd Python3 テキストリンク

目的

物体認識のデータセットにおいてファイル名を適当なものにしておくと、どの画像にどんなオブジェクトがあるのか分からず、データセットの管理が困難になる。
また、ラベリング後に物体認識データセットの名前を変える際は、画像名とアノテーションファイル名、アノテーションファイルの中身の３つを変更する必要がある。特に、アノテーションファイルの中身は簡単には変更することができない。
今回は、これらの3つを一気に変更するプログラムを作成する。

　　　　　　　アノテーションファイル　　　　　　　　　　画像
　　　　　　　　　　　　　　

アノテーションファイルの中身

使用したもの

ラベリング用ソフト

物体認識アルゴリズム

[YOLO]https://github.com/qqwweee/keras-yolo3
[SSD]https://github.com/rykov8/ssd_keras

処理の流れ

  1. python rename_img_ann.py　を実行
  2. "Please input file path"　のあとに、対象データセットへのパスを入力
  3. "Please input new name" のあとに、新しいファイル名を入力する

以下は、../dataset内に存在するファイルをobjectという名前に変更した例である。

結果

アノテーションファイル名、画像名、アノテーションファイルの中身の3つを同時に変更できた。

　　　　　　　　　アノテーションファイル　　　　　　　　画像
　　　　　　　　　　　　　　　　　　　

アノテーションファイルの中身

結論

ラベリング後のファイル名変更といった面倒な作業を自動化できた。
また、今回は例としてobjectという名前に変更したが、dog, catのような名前に変更することで一目でデータセットを管理できるようになる。

プログラム

プログラムはgithub上に公開済　https://github.com/ptxyasu/rename_img_ann
物体認識のデータセットの形式に合わせてあり、JPEGImagesとAnnotationsに存在する同一名ファイル（拡張子を除いた）のファイル名を入力したものに変更している。また、アノテーションファイルの中身はpathタグとfilenameタグのみを変更している。
同一名のものが存在しない場合、プログラムが終了することに注意

rename_img_ann.py
import os
import argparse
import xml.etree.ElementTree as ET

class rename(object):

    def __init__(self,root,name):
        self.root = root
        self.name = name
        self.Annotations = os.path.join(root,'Annotations')
        self.JPEGImages = os.path.join(root,'JPEGImages')

    def get_files(self):
        #ファイル一覧取得
        file_list_ = os.listdir(self.Annotations)
        #拡張子を抜いてファイル一覧作成
        file_list = list(map(lambda x: os.path.splitext(x)[0],file_list_))
        file_list.sort()
        #jpgとxmlを対応づけたまま名前変更
        for i in range (len(file_list)):
            file_name = self.name + str(i)
            img_path_ = self.JPEGImages + "/" + file_list[i]+".jpg"
            img_path = self.JPEGImages + "/" + file_name+".jpg"
            ann_path_ = self.Annotations + "/" + file_list[i]+".xml"
            ann_path = self.Annotations + "/" + file_name+".xml"

            if img_path_ != img_path:
                os.rename(img_path_,img_path)
            if ann_path_ != ann_path:
                os.rename(ann_path_,ann_path)

            tree = ET.parse(ann_path)
            tag_path = tree.find("path")
            tag_path.text = img_path
            tag_name = tree.find("filename")
            tag_name.text = file_name + ".jpg"

            tree.write(ann_path,"utf-8",True)

def main():
    path = input("Please input file path : ")
    name = input("please input new name : ")
    rn = rename(path,name)
    rn.get_files()

if __name__ == '__main__':
    main()