えいあーるれいの技術日記

ROS2やM5 Stack、Ubuntuについて書いています

Darknet用のtextファイルを使ってトリミング

 トリミングにDarknet用のテキストデータを用いるというくだらないプログラムを作成しました。 github.com

 ただし、このプログラムにはしっかりとした目的があります。私は、定点撮影をした動画から連番画像を生成して特定位置の画像を取得したいと思っていましたが、2000枚を超える画像から特定箇所(2か所)をわざわざGIMPなどのツールを使ってトリミングしたくないと思いました。しかし、convertコマンドもなかなか指定が面倒…

 そこで、convertコマンドの引数を簡単に指定できるようにラベリングした複数個所を自動で切り抜くプログラムを作成しました。

使用方法

使用方法は簡単で、対象のフォルダにある1枚の画像に対してアノテーション作業を行い、

python3 crop_img.py --input_txt <target picture's text path(for darknet)>

のコマンドを打ち込むだけです。OpenCVで切り抜きをするわけではないですが、OpenCVが必要です。Pillowとかを使えばよかったかな…

 darknetに使用されるテキストファイルは、ラベル番号→xの中心位置→yの中心位置→幅→高さの順に記載されており、全て比率で書かれています。(つまり、0~1の範囲)

 このフォーマット意外と応用利くかも!

import subprocess
import sys
import argparse
import warnings
import cv2
import os

warnings.simplefilter('ignore')

def img_save(txt, sample_img_path, output_path, save_count):

    ## get abs path
    abs_path = os.path.dirname(os.path.abspath(sample_img_path))
    output_path = abs_path+"/"+output_path

    sample_img = cv2.imread(sample_img_path)
    img_width  = sample_img.shape[1]
    img_height = sample_img.shape[0]

    width    = int(float(txt[3])*img_width)
    height   = int(float(txt[4])*img_height)
    
    cut_x = int((img_width*float(txt[1])) - width /2)
    cut_y = int((img_height*float(txt[2])) - height/2)

    os.makedirs(output_path+str(save_count), exist_ok=True)

    exec_str="for i in "+abs_path+"/*.jpg;do convert -crop "+str(width)+"x"+str(height)+"+"+str(cut_x)+"+"+str(cut_y)+" ${i} "+output_path+str(save_count)+"/$(basename ${i%.jpg})_"+str(save_count)+".jpg;done"
    os.system(exec_str)


## arg------------------------------------------
parser = argparse.ArgumentParser(description="image generator")
parser.add_argument('--input_txt',default=None)
parser.add_argument('--output_folder',default='data')
parser.add_argument('--format',default='jpg')

args = parser.parse_args()

## jpg name
sample_img_txt = os.path.splitext(args.input_txt)[0]+"."+args.format

read_txt_file = open(args.input_txt,'r')

lines = read_txt_file.readlines()

count = 0
for line in lines:
    txt_str = line.split()
    img_save(txt_str, sample_img_txt, args.output_folder, count)
    count = count + 1
    txt_str.clear()

read_txt_file.close()