ChatGPTに多くの写真から動物を抽出するスクリプトを作ってもらった

ChatGPTに以下のお願いをしたら、ほぼ完成形のスクリプトを作ってくれました。

Linux環境でフォルダ内のJPEG画像の中から動物の部分を切り出してファイル名に_数字をつけて出力するシェルスクリプト

追加で切り出した画像は正方形で余白は白ではなく元画像を１．５倍の大きめで切り取りをするよう追加で指示をしたら、以下のスクリプトになりました。（実際にはJPEGのファイル名を小文字ではなく大文字に手作業で修正しています）

まず、シェルスクリプト process_image.sh です。

#!/bin/bash

# 入力フォルダと出力フォルダの指定
INPUT_DIR="./input"
OUTPUT_DIR="./output"

# 出力フォルダが存在しない場合は作成
mkdir -p "$OUTPUT_DIR"

# ファイル番号カウンタ
COUNTER=1

# フォルダ内のJPEG画像を処理
for IMAGE in "$INPUT_DIR"/*.JPG; do
    if [ -f "$IMAGE" ]; then
        # Pythonスクリプトを呼び出して処理
        python3 detect_animals.py "$IMAGE" "$OUTPUT_DIR" "$COUNTER"
        COUNTER=$((COUNTER + 1))
    fi
done

echo "処理が完了しました。出力は $OUTPUT_DIR に保存されています。"

次に、Pythonのスクリプト detect_animals.py です。

import sys
import os
import cv2
import torch
from pathlib import Path

def detect_animals(image_path, output_dir, counter):
    # YOLOv5のモデルをロード
    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

    # 画像を読み込む
    img = cv2.imread(image_path)
    if img is None:
        print(f"Error: Unable to load image {image_path}")
        return

    img_h, img_w, _ = img.shape  # 元画像の高さと幅を取得

    # 推論を実行
    results = model(image_path)

    # 検出されたオブジェクトをループ処理
    detections = results.xyxy[0].numpy()
    for i, det in enumerate(detections):
        # クラスラベルを取得（'animal'クラスを仮定）
        class_id = int(det[5])
        label = model.names[class_id]
        if label in ["dog", "cat", "bird", "horse", "cow"]:  # 動物クラスを指定
            x1, y1, x2, y2 = map(int, det[:4])  # バウンディングボックス座標

            # バウンディングボックスの中心と幅・高さを計算
            cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
            box_w, box_h = x2 - x1, y2 - y1

            # 正方形サイズを計算（バウンディングボックスの長い方を1.5倍に拡大）
            square_size = int(max(box_w, box_h) * 1.5)

            # 正方形範囲を計算（拡大後の範囲を中心基準で計算）
            new_x1 = max(0, cx - square_size // 2)
            new_y1 = max(0, cy - square_size // 2)
            new_x2 = min(img_w, cx + square_size // 2)
            new_y2 = min(img_h, cy + square_size // 2)

            # 元画像から切り出し
            cropped = img[new_y1:new_y2, new_x1:new_x2]

            # 保存ファイル名の作成
            input_name = Path(image_path).stem
            output_path = os.path.join(output_dir, f"{input_name}_{counter}_{i+1}.jpg")

            # 切り出した画像を保存
            cv2.imwrite(output_path, cropped)
            print(f"Saved: {output_path}")

if __name__ == "__main__":
    # コマンドライン引数の取得
    if len(sys.argv) != 4:
        print("Usage: python detect_animals.py <image_path> <output_dir> <counter>")
        sys.exit(1)

    image_path = sys.argv[1]
    output_dir = sys.argv[2]
    counter = int(sys.argv[3])

    detect_animals(image_path, output_dir, counter)

環境構築は以下の手順です。

cd python/
python3 -m venv cropbird
cd cropbird/
source bin/activate
pip install opencv-python yolov5
bash process_images.sh 

YOLOv5の重みファイル（yolov5s.ptなど）を探してきて ~/python/cropbird に置いて、さらに~/python/cropbird/input に処理対象の画像ファイルを置いて、最後の行を実行するとあっさり処理が完了。自力で調べるとかなり時間がかかりますが、あっという間にできてしまいました。ChatGPT（LLM）の威力は凄いです。

できた画像にはいくつか動物が写っていないものなどがありましたので、それを削除してから結合したものがこれです。

十分な成果だと思います。