犬ストレージ・猫ストレージを作る方法（yolo）

初めに

自分の撮りためてた写真の中に、飼っている犬・猫の写真が紛れ込んでいます。家族に写真を共有したいと思いますが、いちいち犬の写真だけ手作業でコピーしていくのは大変です。犬の写真だけ自動で取り出したフォルダがあれば、とても便利だったことでしょう！

犬だけ猫だけのＵＳＢを作ってしまい、犬ストレージ・猫ストレージを作成しましょう！

pythonのtorchを使ったYOLO ver5を使用します。本来YOLOは、写真中の物体検知するモデルで、物体の位置とその物体が何であるかを可能性を教えてくれます。自動運転等で研究が盛んにおこなわれているようです。たとえば、YOLOを使って犬の写真で物体検知すると、以下のような結果が得られます。

すでに学習させてあるモデルを使用する場合には、物体検知できるclassは80個です。猫は15番目、犬は16番目です。

['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 
        'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 
        'teddy bear', 'hair drier', 'toothbrush']

YOLOを使ったコード基本形

さきほどの犬１匹の写真を読み込ませます。

import torch

model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
dir = '/content/DSC_5456.JPG'
results = model(dir)
print(results.pandas().xyxy[0])

2番目のねこ２匹の画像を読み込ませると以下のような結果が得られます。

つまり、dogストレージ、catストレージを作るためにはこれらの結果をつかってif statementで分類すればいいわけです。

犬ストレージ作成　コード

import torch
import glob
import shutil
import os

model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
# 写真フォルダの指定
image_paths = glob.glob("C:\\Users\\pictures\\*.JPG")
# 犬ストレージの指定
output_folder = "C:\\Users\\Dog_strage\\"

for image in image_paths:
    # Inference
    results = model(image)
    print(results.pandas().xyxy)
    print(type(results.pandas().xyxy[0]))
    if sum(results.pandas().xyxy[0]["class"] == 16) >= 1:
        basename = os.path.basename(image)
        full_path = output_folder + basename
        shutil.copy2(image, full_path)

猫ストレージ作成　コード

import torch
import glob
import shutil
import os

model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
# 写真フォルダの指定
image_paths = glob.glob("C:\\Users\\pictures\\*.JPG")
# 猫ストレージの指定
output_folder = "C:\\Users\\cat_strage\\"

for image in image_paths:
    # Inference
    results = model(image)
    print(results.pandas().xyxy)
    print(type(results.pandas().xyxy[0]))
    if sum(results.pandas().xyxy[0]["class"] == 15) >= 1:
        basename = os.path.basename(image)
        full_path = output_folder + basename
        shutil.copy2(image, full_path)

実際にやってみましたが、うまく犬だけの写真が入ったフォルダができました！！これは使えるぞ