プログラミングの芽

面白そうな技術を使って、楽しいことやりたい

機械学習に使う画像を水増しした分のアノテーションファイルを作成する。

以下の記事の作業を行ったという想定の元進めていきます。
atarms.hatenablog.com
atarms.hatenablog.com

正解ラベルの部分をトリミングしているため、画像全体が正解という事になります。
それを90度、180度、270度と回転させただけなので、画像全体が正解という事に変わりはありません。
従って、すべての画像は画像の解像度=ラベルの位置ということになります。

正解ラベルは、上の記事通りにやっていれば
元画像0~600.jpg、90度回転画像601~1201.jpg、180度回転画像...
といった感じで管理していると思うので、アノテーションファイルのラベルを600個1ユニットとしてコピペしていくだけでできます。

必要なもの

import os
import numpy as np
import re
import cv2
import csv


配列に画像のファイル名を入れる

dir_path="./"
img_path=os.listdir(dir_path)

img_path2 = []
for x in img_path:
    m = re.search("[0-9]+", x)
    tuple = (m.group(), x)
    img_path2.append(tuple)
img_path2.sort(key = lambda x: (int(x[0])))

for x in img_path2:
    img_path = np.append(img_path,x[1])
img_path = [x[1] for x in img_path2]

img_path = np.array(img_path,dtype="unicode")


画像の相対パスを配列に入れる

img_path2=[]
for i in range(img_path.shape[0]):
    img_path2.append(dir_path+img_path[i])
img_path2=np.array(img_path2,dtype="unicode")


画像の解像度を配列に入れる(正解ラベルの座標に使う)

img_size=[]#画像サイズ用
for i in range(img_path2.shape[0]):
    img=cv2.imread(img_path2[i])
    img_size.append(img.shape[0:2])
img_size=np.array(img_size,dtype="int64")


元画像のアノテーションファイルを配列に入れる(正解ラベル抽出のため)

origin_csv=[]
with open('labels.csv') as f:
    reader = csv.reader(f)
    header = next(reader) #headerを読み飛ばす
    for line in f: #1行ずつ読み込む(行数分でfor)
        line = line.rstrip("\n").split(",")
        origin_csv.append(line) #csv_data[]に1行ずつ追加
origin_csv = np.array(origin_csv,dtype="unicode")


アノテーションファイル作成

img_csv=[]

img_csv.insert(0,["filename", "width", "height", "class", "xmin", "ymin", "xmax", "ymax"]) #ヘッダ作成

class_cnt = 0
for i in range(img_path2.shape[0]):
    if origin_csv.shape[0] == class_cnt:
        class_cnt = 0
        img_csv.append([img_path[i], img_size[i,0], img_size[i,1], origin_csv[class_cnt,3], 0, 0, img_size[i,0], img_size[i,1]])
        class_cnt = class_cnt + 1
    else:
        img_csv.append([img_path[i], img_size[i,0], img_size[i,1], origin_csv[class_cnt,3], 0, 0, img_size[i,0], img_size[i,1]])
        class_cnt = class_cnt + 1

img_csv=np.array(img_csv)
with open("new_labels.csv","a",newline="") as file:
    writer = csv.writer(file)
    for i in range(len(img_csv)):
        writer.writerow(img_csv[i])


こんなアノテーションファイルができる。
f:id:atarms:20190130211811p:plain

かなりコードがぐちゃぐちゃしてます。ごめんなさい。
備忘録の面が強いのでご容赦。