【物体検出】OIDv4-ToolKit から作成したデータセットでお手軽に RetinaNet を実装

YOLOv3 論文の性能比較で目立っていた RetinaNet に興味がわきました。

深層学習による物体検出は、物体候補領域提案と、回帰分類を分けて行う two-stage 型と、候補領域提案を行わない one-statge 型があり、one-stage は、two-stage よりも高速な推論ができる反面、精度が劣る傾向にあるとのこと。（最近は BoundingBoxを各隅4点でなく、中心座標やコーナーの２点等で予測する Keypoint系もあるようです）

RetinaNet は損失関数を工夫すること（Focal Loss）で、one-stage で精度向上を図ったモデルだそう。

（YOLOv3 論文より）

今回は、だいぶラップされてとても実装が簡単なレポジトリを発見したので、こちらで Dolphin & Shark 検出器を作りたいと思います！

https://github.com/Tessellate-Imaging/Monk_Object_Detection

データは OIDv4-Tookkit でダウンロードしたものを加工して COCO フォーマットに変換する必要があるようです。

レポジトリにあるコードがそのままでは動かなかった＋ちょっと読みづらかったので、少しだけデータセット作成のところは自分なりにアレンジしてやってみました。

データセット作成

まずは環境合わせです。Google Colab バージョンとそうでないもので 2種類 requirements.txt が用意されています。

!cd Monk_Object_Detection/5_pytorch_retinanet/installation && cat requirements_colab.txt | xargs -n 1 -L 1 pip install

1	!cd Monk_Object_Detection/5_pytorch_retinanet/installation && cat requirements_colab.txt \| xargs -n 1 -L 1 pip install

OIDv4-TookKit でデータを収集しておきます。そちらに関してこちらで。

【YOLOv3 データセット作成の旅①】Open Images Dataset の特定クラスの画像のみダウンロード

OIDv4-ToolKit という名前のフォルダをそのまま使ってデータ作成をします。ぱっとやれるように関数にまとめました。

import os
import sys
import cv2
import json
from tqdm.notebook import tqdm
import pandas as pd

def make_dataset(root=None, type="train"):
    # パスの定義
    images_folder = root + type + "/"

    print("-------------------------------------")
    print("making train_labels.csv")
    # 複数クラスの画像をひとまとめにする
    combined = []
    # クラスリストを作成
    complete_list = sorted(os.listdir(images_folder))
    # complete_list に DS_Store が含まれるため削除
    complete_list = complete_list[1:]
    print(complete_list)

    for i in tqdm(range(len(complete_list))):
        name = complete_list[i]

        # アノテーションデータが格納されているフォルダの一覧を取得
        files = os.listdir(images_folder + name + "/Label/")

        for i in tqdm(range(len(files))):
            f = open(images_folder + name + "/Label/" + files[i])
            # ファイル名
            img_name = files[i].split(".")[0] + ".jpg"
            # 画像を読み込み
            img = cv2.imread(images_folder + name + "/" + img_name)
            # アノテーションデータの各行を読み込み
            lines = f.readlines()
            f.close()

            anno = ""

            for j in range(len(lines)):
                # 各座標情報を取得
                tmp = lines[j].split(" ")
                x1 = int(float(tmp[-4]))
                y1 = int(float(tmp[-3]))
                x2 = int(float(tmp[-2]))
                y2 = int(float(tmp[-1]))

                anno += str(x1) + " " + str(y1) + " " + str(x2) + " " + str(y2) + " " + name + " "
                
            anno = anno[:-1]
            combined.append([name + "/" + img_name, anno])

    df = pd.DataFrame(combined, columns=["ID", "Labels"])
    df.to_csv(images_folder + "train_labels.csv", index=False)

    print("done")

    print("-------------------------------------")
    print("making classes.txt")

    annotations_path = root + "/annotations/"

    if not os.path.isdir(annotations_path):
        os.mkdir(annotations_path)
        
    input_annotations_path = images_folder + "/train_labels.csv"
    output_annotation_file = annotations_path + "/instances_" + type + ".json"
    output_classes_file = annotations_path + "/classes.txt"

    # annotation フォルダがなければ作成
    if not os.path.isdir(annotations_path):
        os.mkdir(annotations_path);

    # train_labels.csv
    df = pd.read_csv(input_annotations_path)
    columns = df.columns
    delimiter = " "

    anno = []

    for i in range(len(df)):
        img_name = df[columns[0]][i]
        labels = df[columns[1]][i]
        # classesを格納したリスト
        tmp = labels.split(delimiter)
        for j in range(len(tmp) // 5):
            label = tmp[j*5+4]
            if (label not in anno):
                anno.append(label)
        
    anno = sorted(anno)

    list_dict = []

    for i in tqdm(range(len(anno))):
        tmp = {}
        tmp["supercategory"] = "master"
        tmp["id"] = i
        tmp["name"] = anno[i]
        list_dict.append(tmp)

    # classes.txt を作成
    anno_f = open(output_classes_file, "w")
    for i in range(len(anno)):
        anno_f.write(anno[i] + "¥n")
    anno_f.close()

    print("done")

    print("-------------------------------------")
    print("making isinstance_train.json")

    # COCOフォーマットのデータ作成
    coco_data = {}
    coco_data["type"] = "instances"
    coco_data["images"] = []
    coco_data["annotations"] = []
    coco_data["categories"] = list_dict

    image_id = 0
    annotation_id = 0

    for i in tqdm(range(len(df))):
        img_name = df[columns[0]][i]
        labels = df[columns[1]][i]
        tmp = labels.split(delimiter)
        image_in_path = images_folder + "/" + img_name
        img = cv2.imread(image_in_path, 1)
        h, w, c = img.shape

        images_tmp = {}
        images_tmp["file_name"] = img_name
        images_tmp["height"] = h
        images_tmp["width"] = w
        images_tmp["id"] = image_id
        coco_data["images"].append(images_tmp)

        for j in range(len(tmp)//5):
            x1 = int(tmp[j*5+0])
            y1 = int(tmp[j*5+1])
            x2 = int(tmp[j*5+2])
            y2 = int(tmp[j*5+3])
            label = tmp[j*5+4]
            annotations_tmp = {}
            annotations_tmp["id"] = annotation_id
            annotation_id += 1
            annotations_tmp["image_id"] = image_id
            annotations_tmp["segmentation"] = []
            annotations_tmp["ignore"] = 0
            annotations_tmp["area"] = (x2-x1)*(y2-y1)
            annotations_tmp["iscrowd"] = 0
            annotations_tmp["bbox"] = [x1, y1, x2-x1, y2-y1]
            annotations_tmp["category_id"] = anno.index(label)

            coco_data["annotations"].append(annotations_tmp)
        image_id += 1;
  
    # isinstance_train.json を作成
    outfile =  open(output_annotation_file, 'w');
    json_str = json.dumps(coco_data, indent=4);
    outfile.write(json_str);
    outfile.close();

    print("done")

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

import os

import sys

import cv2

import json

from tqdm.notebook import tqdm

import pandas as pd

def make_dataset(root=None, type="train"):

# パスの定義

images_folder = root + type + "/"

print("-------------------------------------")

print("making train_labels.csv")

# 複数クラスの画像をひとまとめにする

combined = []

# クラスリストを作成

complete_list = sorted(os.listdir(images_folder))

# complete_list に DS_Store が含まれるため削除

complete_list = complete_list[1:]

print(complete_list)

for i in tqdm(range(len(complete_list))):

name = complete_list[i]

# アノテーションデータが格納されているフォルダの一覧を取得

files = os.listdir(images_folder + name + "/Label/")

for i in tqdm(range(len(files))):

f = open(images_folder + name + "/Label/" + files[i])

# ファイル名

img_name = files[i].split(".")[0] + ".jpg"

# 画像を読み込み

img = cv2.imread(images_folder + name + "/" + img_name)

# アノテーションデータの各行を読み込み

lines = f.readlines()

f.close()

anno = ""

for j in range(len(lines)):

# 各座標情報を取得

tmp = lines[j].split(" ")

x1 = int(float(tmp[-4]))

y1 = int(float(tmp[-3]))

x2 = int(float(tmp[-2]))

y2 = int(float(tmp[-1]))

anno += str(x1) + " " + str(y1) + " " + str(x2) + " " + str(y2) + " " + name + " "

anno = anno[:-1]

combined.append([name + "/" + img_name, anno])

df = pd.DataFrame(combined, columns=["ID", "Labels"])

df.to_csv(images_folder + "train_labels.csv", index=False)

print("done")

print("-------------------------------------")

print("making classes.txt")

annotations_path = root + "/annotations/"

if not os.path.isdir(annotations_path):

os.mkdir(annotations_path)

input_annotations_path = images_folder + "/train_labels.csv"

output_annotation_file = annotations_path + "/instances_" + type + ".json"

output_classes_file = annotations_path + "/classes.txt"

# annotation フォルダがなければ作成

if not os.path.isdir(annotations_path):

os.mkdir(annotations_path);

# train_labels.csv

df = pd.read_csv(input_annotations_path)

columns = df.columns

delimiter = " "

anno = []

for i in range(len(df)):

img_name = df[columns[0]][i]

labels = df[columns[1]][i]

# classesを格納したリスト

tmp = labels.split(delimiter)

for j in range(len(tmp) // 5):

label = tmp[j*5+4]

if (label not in anno):

anno.append(label)

anno = sorted(anno)

list_dict = []

for i in tqdm(range(len(anno))):

tmp = {}

tmp["supercategory"] = "master"

tmp["id"] = i

tmp["name"] = anno[i]

list_dict.append(tmp)

# classes.txt を作成

anno_f = open(output_classes_file, "w")

for i in range(len(anno)):

anno_f.write(anno[i] + "¥n")

anno_f.close()

print("done")

print("-------------------------------------")

print("making isinstance_train.json")

# COCOフォーマットのデータ作成

coco_data = {}

coco_data["type"] = "instances"

coco_data["images"] = []

coco_data["annotations"] = []

coco_data["categories"] = list_dict

image_id = 0

annotation_id = 0

for i in tqdm(range(len(df))):

img_name = df[columns[0]][i]

labels = df[columns[1]][i]

tmp = labels.split(delimiter)

image_in_path = images_folder + "/" + img_name

img = cv2.imread(image_in_path, 1)

h, w, c = img.shape

images_tmp = {}

images_tmp["file_name"] = img_name

images_tmp["height"] = h

images_tmp["width"] = w

images_tmp["id"] = image_id

coco_data["images"].append(images_tmp)

for j in range(len(tmp)//5):

x1 = int(tmp[j*5+0])

y1 = int(tmp[j*5+1])

x2 = int(tmp[j*5+2])

y2 = int(tmp[j*5+3])

label = tmp[j*5+4]

annotations_tmp = {}

annotations_tmp["id"] = annotation_id

annotation_id += 1

annotations_tmp["image_id"] = image_id

annotations_tmp["segmentation"] = []

annotations_tmp["ignore"] = 0

annotations_tmp["area"] = (x2-x1)*(y2-y1)

annotations_tmp["iscrowd"] = 0

annotations_tmp["bbox"] = [x1, y1, x2-x1, y2-y1]

annotations_tmp["category_id"] = anno.index(label)

coco_data["annotations"].append(annotations_tmp)

image_id += 1;

# isinstance_train.json を作成

outfile = open(output_annotation_file, 'w');

json_str = json.dumps(coco_data, indent=4);

outfile.write(json_str);

outfile.close();

print("done")

root = "/OIDv4_ToolKit/OID/Dataset/"

make_dataset(root, "train")

make_dataset(root, "validation")

make_dataset(root, "test")

root = "/OIDv4_ToolKit/OID/Dataset/"

make_dataset(root, "train")

make_dataset(root, "validation")

make_dataset(root, "test")

RetinaNet の学習

データセットが作成できたため、学習に取り掛かります。ここからはこのレポジトリは、少ないコードで進められて便利。

まずは import できるようにパスを通してあげます。

import os 
import sys

sys.path.append("/Monk_Object_Detection/5_pytorch_retinanet/lib/")

from train_detector import Detector

import os

import sys

sys.path.append("/Monk_Object_Detection/5_pytorch_retinanet/lib/")

from train_detector import Detector

訓練、検証DataLoader を作成します。

detector = Detector()

# 訓練用
root_dir = "/OIDv4_ToolKit/OID/Dataset"
coco_dir = ""
img_dir = ""
set_dir = "train"

detector.Train_Dataset(root_dir, coco_dir, img_dir, set_dir, batch_size=8, use_gpu=True)

# 検証用
root_dir = "/OIDv4_ToolKit/OID/Dataset"
coco_dir = ""
img_dir = ""
set_dir = "validation"

detector.Val_Dataset(root_dir, coco_dir, img_dir, set_dir)

detector = Detector()

# 訓練用

root_dir = "/OIDv4_ToolKit/OID/Dataset"

coco_dir = ""

img_dir = ""

set_dir = "train"

detector.Train_Dataset(root_dir, coco_dir, img_dir, set_dir, batch_size=8, use_gpu=True)

# 検証用

root_dir = "/OIDv4_ToolKit/OID/Dataset"

coco_dir = ""

img_dir = ""

set_dir = "validation"

detector.Val_Dataset(root_dir, coco_dir, img_dir, set_dir)

バックボーンには ResNet18 / 34/ 50/ 101 が選べます。

detector.Model(model_name="resnet34")

detector.Set_Hyperparams(lr=1e-5, val_interval=1, print_interval=10)

detector.Train(num_epochs=5, output_model_name="model.pt")

detector.Model(model_name="resnet34")

detector.Set_Hyperparams(lr=1e-5, val_interval=1, print_interval=10)

detector.Train(num_epochs=5, output_model_name="model.pt")

最適化手法は Adam となっているため、train_detector.py の中身を変更して SGD 等にしても良いと思います。

def Set_Hyperparams(self, lr=0.0001, val_interval=1, print_interval=20):
        self.system_dict["params"]["lr"] = lr;
        self.system_dict["params"]["val_interval"] = val_interval;
        self.system_dict["params"]["print_interval"] = print_interval;

        self.system_dict["local"]["optimizer"] = torch.optim.Adam(self.system_dict["local"]["model"].parameters(), 
                                                                    self.system_dict["params"]["lr"]);

        self.system_dict["local"]["scheduler"] = torch.optim.lr_scheduler.ReduceLROnPlateau(self.system_dict["local"]["optimizer"], 
                                                                    patience=3, verbose=True)

        self.system_dict["local"]["loss_hist"] = collections.deque(maxlen=500)

def Set_Hyperparams(self, lr=0.0001, val_interval=1, print_interval=20):

self.system_dict["params"]["lr"] = lr;

self.system_dict["params"]["val_interval"] = val_interval;

self.system_dict["params"]["print_interval"] = print_interval;

self.system_dict["local"]["optimizer"] = torch.optim.Adam(self.system_dict["local"]["model"].parameters(),

self.system_dict["params"]["lr"]);

self.system_dict["local"]["scheduler"] = torch.optim.lr_scheduler.ReduceLROnPlateau(self.system_dict["local"]["optimizer"],

patience=3, verbose=True)

self.system_dict["local"]["loss_hist"] = collections.deque(maxlen=500)

ログはこんな形で出てきます。

推論

学習が終わると、pt ファイルが出来上がるので、それを読み込みます。

from infer_detector import Infer

model = Infer()
model.Model(model_path="model.pt")
class_list = ["Dolphin", "shark"]

from infer_detector import Infer

model = Infer()

model.Model(model_path="model.pt")

class_list = ["Dolphin", "shark"]

いくつか推論してみます。

from IPython.display import Image

test_img = "/OIDv4_ToolKit/OID/Dataset/test/shark/2613ccddca33df02.jpg"

scores, labels, boxse = model.Predict(test_img, class_list, vis_threshold=0.5)

Image(filename="output.jpg")