search
LoginSignup
23

4行のコードでAI物体検出を使う方法【Python,Yolov5】

物体検出を使ってみたい

プログラムが自動で画像から物体を見つけてくれる物体検出、色々な分野で使えます。
6ffb75b6-a203-4b7f-b512-504062e15991.jpeg

AIを使うのは難しそう

何行もコードが必要と思うじゃないですか。

実はめちゃくちゃかんたん

すごくかんたんに使えるようになっています。
今回はYolov5を使います。

方法

torch hubからモデルをロードして、画像を入力として与えます。

import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # initialize model
result = model('teddybear.jpg') # inference
result.save()

これだけで推論画像が保存されます。
かんたんでしょ?

6ffb75b6-a203-4b7f-b512-504062e15991.jpeg

結果を見る

結果の座標にアクセスすることもできます。

result.pandas().xywhn

[ xcenter ycenter width height confidence class name
0 0.701028 0.522462 0.262828 0.421278 0.930700 77 teddy bear
1 0.703856 0.719701 0.487835 0.442845 0.913001 13 bench]

検出可能なクラスは以下で取得できます。

model.names

{0: 'person',
1: 'bicycle',
2: 'car',
3: 'motorcycle',
4: 'airplane',
5: 'bus',
6: 'train',
7: 'truck',
8: 'boat',
9: 'traffic light',
10: 'fire hydrant',
11: 'stop sign',
12: 'parking meter',
13: 'bench',
14: 'bird',
15: 'cat',
16: 'dog',
17: 'horse',
18: 'sheep',
19: 'cow',
20: 'elephant',
21: 'bear',
22: 'zebra',
23: 'giraffe',
24: 'backpack',
25: 'umbrella',
26: 'handbag',
27: 'tie',
28: 'suitcase',
29: 'frisbee',
30: 'skis',
31: 'snowboard',
32: 'sports ball',
33: 'kite',
34: 'baseball bat',
35: 'baseball glove',
36: 'skateboard',
37: 'surfboard',
38: 'tennis racket',
39: 'bottle',
40: 'wine glass',
41: 'cup',
42: 'fork',
43: 'knife',
44: 'spoon',
45: 'bowl',
46: 'banana',
47: 'apple',
48: 'sandwich',
49: 'orange',
50: 'broccoli',
51: 'carrot',
52: 'hot dog',
53: 'pizza',
54: 'donut',
55: 'cake',
56: 'chair',
57: 'couch',
58: 'potted plant',
59: 'bed',
60: 'dining table',
61: 'toilet',
62: 'tv',
63: 'laptop',
64: 'mouse',
65: 'remote',
66: 'keyboard',
67: 'cell phone',
68: 'microwave',
69: 'oven',
70: 'toaster',
71: 'sink',
72: 'refrigerator',
73: 'book',
74: 'clock',
75: 'vase',
76: 'scissors',
77: 'teddy bear',
78: 'hair drier',
79: 'toothbrush'}

自分でトレーニングしたモデルの使用

自分でトレーニングしたモデルを使うこともできます。
これにより、デフォルトのクラス以外の物体も検出することができます。

model = torch.hub.load('ultralytics/yolov5', 'custom', path='path/to/best.pt') # Specify your own path

🐣


フリーランスエンジニアです。
お仕事のご相談こちらまで
rockyshikoku@gmail.com

機械学習、ARアプリ(Web/iOS)を作っています。
機械学習/AR関連の情報を発信しています。

Twitter
Medium
GitHub

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
What you can do with signing up
23