csv形式のanotationファイルをpklに変換する

初めに

初投稿です、お手柔らかに。datasetのanotationがcsv形式だったので、Single Shot Multibox Detector(SSD)で使えるようにpklファイルに変換するときに書いたプログラムを置いておきます。

ちなみに、SSDの参考にさせていただいたのはこちらの記事です。

https://qiita.com/slowsingle/items/64cc927bb29a49a7af14

数年前の記事なので、最新のverを使っている人はいくつか修正が必要だと思います。

プログラム

日本語のコメントをうざったいぐらい入れてます（笑）

csv_for_pkl.py

import pickle as pkl
import os
import csv
import numpy as np

class CSV_preprocessor(object):
    def __init__(self, data_path):
        self.path = data_path
        self.num_classes = 2   #クラス数
        self.data = dict()
        self._preprocess_CSV()

    def _preprocess_CSV(self):
        with open(self.path, 'r') as f: #fileをfとして読み込み
            reader = csv.reader(f)
            width =900   #画像の幅、anotationの中にあるならfor文内で取得してください
            height = 600   #画像の高さ
            flag = 1   #1行目をスキップするためのflag
            bounding_boxes = []
            one_hot_classes = []
            af = 'Non'

            for row in reader:
                if row[1] != af and af != 'Non':   #前の行のimage_idと異なるなら保存処理
                    image_name = af + '.jpg'
                    bounding_boxes = np.asarray(bounding_boxes)
                    one_hot_classes = np.asarray(one_hot_classes)
                    image_data = np.hstack((bounding_boxes, one_hot_classes))
                    self.data[image_name] = image_data
                    bounding_boxes =[]
                    one_hot_classes =[]

                if flag == 1:
                    flag += 1
                    continue
                #anotationがboundingboxの４点を持っていたので以下のような書き方に
                #なりました.値の場所が不変ならその値を入力してください
                xmin = (min([float(row[2]), float(row[4]), float(row[6]), float(row[8])])) / width
                ymin = (min([float(row[3]), float(row[5]), float(row[7]), float(row[9])])) / height
                xmax = (max([float(row[2]), float(row[4]), float(row[6]), float(row[8])])) / width
                ymax = (max([float(row[3]), float(row[5]), float(row[7]), float(row[9])])) / height
                bounding_box = [xmin, ymin, xmax, ymax]
                bounding_boxes.append(bounding_box)
                class_name = row[10]   #classの名前の場所
                one_hot_class = self._to_one_hot(class_name)
                one_hot_classes.append(one_hot_class)

                af = row[1]   #画像の名前の場所
 
            image_name = af + '.jpg'   #jpgで保存します
            bounding_boxes = np.asarray(bounding_boxes)
            one_hot_classes = np.asarray(one_hot_classes)
            image_data = np.hstack((bounding_boxes, one_hot_classes))
            self.data[image_name] = image_data
            bounding_boxes =[]
            one_hot_classes =[]

    def _to_one_hot(self, name):
        one_hot_vector = [0] * self.num_classes
        if name == 'class name':
            one_hot_vector[0] = 1
        
        elif name == 'class name':
            one_hot_vector[1] =1

        else :
            print('unknown label: %s' %name)

        return one_hot_vector

data = CSV_preprocessor('anotation pass').data
pkl.dump(data, open('train.pkl','wb'))

走り書きで書いたので、たぶんいろいろ直せるところはあると思います。
良かったらコメントしていってください。