こんにちは。現役エンジニアの三年坊主（@SannenBouzu）です。
エンジニア歴は6年。うち2年以上のAI開発を経験し、Web技術とAI技術をバランスよく習得してきました。
今回は、PyTorchの学習済みモデルを実行し、画像から物体検出を行う方法を紹介します。
（PyTorchの基礎、物体検出の基礎については、別記事で取り扱う予定です）
「学習済みモデル」を使うことで、「データ収集」「前処理」「学習」といった、ディープラーニングで時間がかかる多くのプロセスをスキップして、「モデルを使った推論」という結果だけを素早く得ることができます。
また、記事で扱う「モデルを定義→入力を与える→出力を得る」という「推論」の流れは、物体検出に限りません。一度理解しておくと、ディープラーニング、もっと広く機械学習全般にも応用が可能です。
記事のまとめ部分では、全体像を整理しやすいように、ObjectDetectorというクラスを定義していますので、復習に役立ててください。
- PyTorchを使った物体検出をサクッと試してみたい方
- Pythonの物体検出プログラムを確認したい方
ライブラリのインポート（PyTorchや画像処理ライブラリなど）
（環境の準備については、別記事で取り扱う予定です。Google Colaboratory、またはDockerでの環境構築をおすすめします）
PyTorchなど、今回使うライブラリをインポートします。
- PyTorch：Pythonのオープンソースの機械学習ライブラリ
- torchvision：PyTorchのパッケージ。コンピュータビジョンにおける有名なデータセット・モデルアーキテクチャ・画像変換処理などから構成される
- Pillow(PIL)：Pythonの画像処理ライブラリ
- OpenCV(cv2)：Pythonの画像処理ライブラリ
1 2 3 4 5 6 7 8 9 | import cv2 import torch import torchvision import matplotlib.pyplot as plt import numpy as np from PIL import Image #, ImageDraw, ImageFont from torchvision.models.detection.faster_rcnn import FastRCNNPredictor |
今回の記事は、こちらのバージョンで動作確認しています。
1 2 3 4 5 6 7 | print(torch.__version__) print(torchvision.__version__) print(cv2.__version__) 1.6.0 0.7.0 4.4.0 |
GPU/CPUの確認と切り替え
PyTorchといえばディープラーニング。ディープラーニングといえば、主に画像処理を高速化するためにGPU (Graphics Processing Unit) を使う場面が増えています。
GPUを使えるかどうか確認しておきましょう。ちなみに、今回はGPUを使えなくても問題ありません（私も手元のMacで試しました）。
1 2 3 | torch.cuda.is_available() False |
torch.deviceを使って、PyTorchのテンソルをどのデバイスに割り当てるのか指定します（CPUかGPUか、GPUが複数あるときにはどのGPUか）。
1 2 3 4 | device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') print(device) cpu |
PyTorchのインポートを確認する
順番が入れ変わってしまいましたが、適当なテンソルを作って、PyTorchのインポートができているか確認しましょう。
1 2 3 4 5 6 7 8 | x = torch.rand(5, 3) print(x) tensor([[0.5756, 0.4099, 0.4064], [0.0074, 0.1329, 0.7557], [0.6695, 0.4446, 0.3618], [0.8736, 0.7798, 0.6291], [0.4421, 0.5442, 0.5749]]) |
物体検出に使う画像とモデルを用意する
物体検出をするため、適当な画像とモデルを用意します。
画像
飛行機と時計の画像を使ってみます。
1 2 3 4 5 | image1_file_path = './airplane.jpg' image2_file_path = './clock.jpg' img1 = Image.open(image1_file_path).convert("RGB") img2 = Image.open(image2_file_path).convert("RGB") |
学習済みのモデル
torchvisionは、コンピュータビジョンのタスクに便利なデータセット、モデル、画像変換処理などをまとめたパッケージです。
torchvisionには、Faster R-CNNという物体検出モデルが用意されているので、使ってみましょう。
（物体検出の基礎については、別記事で取り扱う予定です）
pretrained=Trueにして、学習済みのモデルを利用します。
1 2 3 4 | def get_object_detection_model(): # load an object detection model pre-trained on COCO model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True) return model |
モデルをdevice(CPUまたはGPU)に割り当てます。
1 2 3 4 5 | # get the model using our helper function model = get_object_detection_model() # move model to the right device model.to(device) |
PIL画像をテンソルに変換する
PILライブラリで読み込んだ画像を、PyTorchの物体検出モデルで使えるよう、テンソルに変換する必要があります。
1 2 3 4 5 6 | transform = torchvision.transforms.Compose([ torchvision.transforms.ToTensor() ]) img1 = transform(img1) img2 = transform(img2) |
img1, img2のsizeを確認しましょう。
1 2 3 4 5 | img1.size() img2.size() torch.Size([3, 1982, 3027]) torch.Size([3, 2560, 3840]) |
物体検出モデルを実行する（予測・推論）
公式ドキュメントで、Faster R-CNNの入力を確認します。
The input to the model is expected to be a list of tensors, each of shape [C, H, W], one for each image, and should be in 0-1 range. Different images can have different sizes.
先ほどテンソルに変換したimg1とimg2のサイズは、
- (3, 1982, 3027)
- (3, 2560, 3840)
どちらも、C（チャネル）, H（高さ）, W（幅）の順になっていましたね。
値の範囲も、0から1の間に収まっています。
1 2 3 4 5 | torch.min(img1) torch.max(img1) tensor(0.) tensor(1.) |
モデルに入力する画像（テンソル）には問題なさそうなので、さっそくモデルに入力してみましょう。
- model.eval()：モデルをevaluationモードに切り替える
- →学習時特有の処理を止める
- torch.no_grad()：自動勾配の計算を止める
- →メモリ使用量を減らせる
- 画像のテンソルは、modelの引数にリストで与える
- →predictionsは画像ごとの物体検出結果を含むリスト
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | score_threshold = 0.9 model.eval() with torch.no_grad(): predictions = model([ img1.to(device), img2.to(device) ]) # filter out predictions with low scores _predictions = [] for pred in predictions: mask = pred['scores'] >= score_threshold _predictions.append( { 'boxes': pred['boxes'][mask], 'labels': pred['labels'][mask], 'scores': pred['scores'][mask], } ) predictions = _predictions |
物体検出モデルの予測結果には、「スコア」が含まれています。confidenceという呼び方もありますが、要するに「どの程度自信を持てる予測結果か」という数字を、0から1の間で出してくれます。
スコアが低い予測結果を受け入れる場合、より多くの物体を検出できる可能性が上がる一方で、正しくない物体検出が増えてしまうかもしれません。
上記のコードでは、score_thresholdに設定した値より低いスコアの予測結果をはじいています。
predictionsに格納した、予測結果を見てみましょう。
1 2 3 4 5 6 7 8 | predictions [{'boxes': tensor([[ 630.8052, 678.9952, 2871.0374, 1278.1471]]), 'labels': tensor([5]), 'scores': tensor([0.9994])}, {'boxes': tensor([[1073.0975, 475.4993, 2837.7283, 2091.3835]]), 'labels': tensor([85]), 'scores': tensor([0.9990])}] |
公式ドキュメントで、Faster R-CNNの出力を確認します。
- boxes (FloatTensor[N, 4]): the predicted boxes in [x1, y1, x2, y2] format, with values of x between 0 and W and values of y between 0 and H
- labels (Int64Tensor[N]): the predicted labels for each image
- scores (Tensor[N]): the scores or each prediction
boxesというのが、バウンディングボックス(bounding box)、物体をちょうど囲むのに必要な大きさの長方形（矩形）のことで、モデルが予測する物体の位置を表しています。
画像を入力すると、画像のどの位置に、何が写っているのか予測する「物体検出」を実行することができました。
バウンディングボックス(物体の位置)を可視化する
物体検出ができたので、人間の目にも分かりやすい形で可視化してみましょう。
さて、今回使った学習済みモデルは、COCO Datasetというデータセットを使って「事前に学習した物体」を検出できるものです。
「事前に学習した物体」というのは、例えば
- 人物
- 自転車
- 車
- オートバイ
- 飛行機
など。具体的には以下の通りで、
1 2 3 4 5 6 7 8 9 10 11 12 13 14 | COCO_INSTANCE_CATEGORY_NAMES = [ '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush' ] |
N/Aと__background__を除くと80種類になっています。
1 2 3 | len(COCO_INSTANCE_CATEGORY_NAMES) 91 |
predictionsのlabelsに
- ‘labels’: tensor([5])
- ‘labels’: tensor([85])
と出ていたのは、画像から予想できる通り、それぞれ「飛行機」と「時計」でした。
1 2 3 4 5 | # 5: airplane, 85: clock print(COCO_INSTANCE_CATEGORY_NAMES[5], COCO_INSTANCE_CATEGORY_NAMES[85]) airplane clock |
cv2.rectangle
続いて、物体の位置を可視化していきましょう。
いくつか方法はありますが、こちらはOpenCVのrectangle関数を使って長方形を描画する方法です。
- OpenCVの関数を使うので、画像もOpenCVで読み込んでいる
- cv2.rectangleは座標をintで指定しないといけないので、predictionsのboxesをintに変換している
- cv2.putTextでラベルの名前（airplane, clockなど）を描画
1 2 3 4 5 | font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 2 color = (255, 0, 0) thickness = 3 delta = 20 |
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 | # OpenCV's rectangle function to overlay bounding boxes on image. # https://stackoverflow.com/a/60275224 inputs = [cv2.imread(path) for path in [image1_file_path, image2_file_path]] for i,img in enumerate(inputs): for (_box, _label, _score) in zip(predictions[i]['boxes'], predictions[i]['labels'], predictions[i]['scores']): # get prediction results: bounding box, label(_name), prediction score box = {k:int(v) for (k,v) in zip(['x0', 'y0', 'x1', 'y1'], _box.tolist())} # If you have a one element tensor, use .item() to get the value as a Python number label = _label.item() label_name = COCO_INSTANCE_CATEGORY_NAMES[label] score = _score.item() # draw bounding box, label and prediction score img = cv2.rectangle(img, (box['x0'], box['y0']), (box['x1'], box['y1']), color, thickness) img = cv2.putText(img, '{} {}'.format(label_name, '{:.1%}'.format(score)), (box['x0'], box['y0']-delta), font, font_scale, color, thickness, cv2.LINE_AA) # display plt.figure(figsize = (12,9)) plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) plt.axis("off") plt.show() |
まとめ: ObjectDetectorクラスを定義して全体像を整理する
今回は、PyTorchの学習済みモデルを実行し、画像から物体検出を行う方法を紹介しました。
最後に、今回の記事で紹介した内容をObjectDetectorというクラスにまとめます。
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 | class ObjectDetector: def __init__(self): self.model = get_object_detection_model() self.paths = None self.predictions = None self.transform = torchvision.transforms.Compose([ torchvision.transforms.ToTensor() ]) # inference mode self.model.eval() def _filter_out_predictions(self, score_threshold=None): # filter out predictions with low scores if score_threshold and 0.0 <= score_threshold <= 1.0: predictions = [] for pred in self.predictions: mask = pred['scores'] >= score_threshold predictions.append( { 'boxes': pred['boxes'][mask], 'labels': pred['labels'][mask], 'scores': pred['scores'][mask], } ) self.predictions = predictions def detect(self, paths: list, score_threshold=None): self.paths = paths # read images as PIL Image, convert PIL images to Tensor inputs = [self.transform(Image.open(path).convert("RGB")) for path in paths] with torch.no_grad(): self.predictions = self.model([ x.to(device) for x in inputs ]) self._filter_out_predictions(score_threshold) return self.predictions def visualize(self): if self.paths is None or self.predictions is None: return # read images as OpenCV image for easier visualization inputs = [cv2.imread(path) for path in self.paths] for i,img in enumerate(inputs): for (_box, _label, _score) in zip(self.predictions[i]['boxes'], self.predictions[i]['labels'], self.predictions[i]['scores']): # box: x0, y0, x1, y1 box = {k:int(v) for (k,v) in zip(['x0', 'y0', 'x1', 'y1'], _box.tolist())} # If you have a one element tensor, use .item() to get the value as a Python number label = _label.item() label_name = COCO_INSTANCE_CATEGORY_NAMES[label] score = _score.item() # draw bounding box, label and prediction score img = cv2.rectangle(img, (box['x0'], box['y0']), (box['x1'], box['y1']), color, thickness) img = cv2.putText(img, '{} {}'.format(label_name, '{:.1%}'.format(score)), (box['x0'], box['y0']-delta), font, font_scale, color, thickness, cv2.LINE_AA) # display plt.figure(figsize = (12,9)) plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) plt.axis("off") plt.show() |
使い方は以下の通り。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | object_detector = ObjectDetector() results = object_detector.detect( [image1_file_path, image2_file_path], score_threshold=0.9 ) [{'boxes': tensor([[ 630.8052, 678.9952, 2871.0374, 1278.1471]]), 'labels': tensor([5]), 'scores': tensor([0.9994])}, {'boxes': tensor([[1073.0975, 475.4993, 2837.7283, 2091.3835]]), 'labels': tensor([85]), 'scores': tensor([0.9990])}] object_detector.visualize() |
「モデルを定義→入力を与える→出力を得る」という「推論」の流れを、実感していただけたかなと思います。
ところで、学習済みモデルを使うと、手軽に物体検出をできる反面、当然ながら、「事前に学習した物体しか検出できない」というデメリットがあります。
（自分で用意したデータ・誰かが用意してくれたデータでモデルを学習させ、好きな物体を検出する方法については、別記事で取り扱う予定です）
