物体検出を使ってみたい
プログラムが自動で画像から物体を見つけてくれる物体検出、色々な分野で使えます。
AIを使うのは難しそう
何行もコードが必要と思うじゃないですか。
実はめちゃくちゃかんたん
すごくかんたんに使えるようになっています。
今回はYolov5を使います。
方法
torch hubからモデルをロードして、画像を入力として与えます。
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # initialize model
result = model('teddybear.jpg') # inference
result.save()
これだけで推論画像が保存されます。
かんたんでしょ？
結果を見る
結果の座標にアクセスすることもできます。
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関連の情報を発信しています。