YOLOv3 論文の性能比較で目立っていた RetinaNet に興味がわきました。
深層学習による物体検出は、物体候補領域提案と、回帰分類を分けて行う two-stage 型と、候補領域提案を行わない one-statge 型があり、one-stage は、two-stage よりも高速な推論ができる反面、精度が劣る傾向にあるとのこと。(最近は BoundingBoxを各隅4点でなく、中心座標やコーナーの2点等で予測する 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 が用意されています。
1 |
!cd Monk_Object_Detection/5_pytorch_retinanet/installation && cat requirements_colab.txt | xargs -n 1 -L 1 pip install |
OIDv4-TookKit でデータを収集しておきます。そちらに関してこちらで。
OIDv4-ToolKit という名前のフォルダをそのまま使ってデータ作成をします。ぱっとやれるように関数にまとめました。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 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") |
1 2 3 4 5 6 7 |
root = "/OIDv4_ToolKit/OID/Dataset/" make_dataset(root, "train") make_dataset(root, "validation") make_dataset(root, "test") |
RetinaNet の学習
データセットが作成できたため、学習に取り掛かります。ここからはこのレポジトリは、少ないコードで進められて便利。
まずは import できるようにパスを通してあげます。
1 2 3 4 5 6 |
import os import sys sys.path.append("/Monk_Object_Detection/5_pytorch_retinanet/lib/") from train_detector import Detector |
訓練、検証DataLoader を作成します。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
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 が選べます。
1 2 3 4 5 |
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 等にしても良いと思います。
1 2 3 4 5 6 7 8 9 10 11 12 |
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 ファイルが出来上がるので、それを読み込みます。
1 2 3 4 5 |
from infer_detector import Infer model = Infer() model.Model(model_path="model.pt") class_list = ["Dolphin", "shark"] |
いくつか推論してみます。
1 2 3 4 5 6 7 |
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") |

最後まで目を通していただきありがとうございました!