YOLOv5のdetect.pyコードを1つずつ検出して保存できるように修正し、各動画のパラメータを個別に操作する

YOLOv5のdetect.pyコードのロジックが本当によくわからない.YOLOv3とYOLOv4のdetectロジックを読んだ後、基本的にopencvを使って各動画を操作している.より明確で理解しやすい.YOLOv5の作者は役に立たないようです. .opencvを操作したり、opencvのビデオ操作を別のpyファイルにカプセル化して隠したりするのはちょっとわかりにくいので、os.listdirを使ってビデオファイルディレクトリ内のすべてのビデオを読み取り、それらを検出するという最も愚かな方法を使用しました一つずつ。同時に、ピクチャーフレームの機能を書き換え(キーフレームの内容を保存するため)、検出コマンドは python detect.py --exist-ok --nosave を使用します。 nosaveのオプションなので浅いです 作者のフレームロジックを見て、opencvのrectangleメソッドが今も使われていることがわかりました(作者の隠し

import numpy as np
import argparse
import os
import sys
from pathlib import Path
import time
import shutil
from PIL import Image
import cv2
import torch
import torch.backends.cudnn as cudnn

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] #
str(ROOT) が sys.path にない場合の YOLOv5 ルート ディレクトリ:
    sys.path.append(str(ROOT)) # ROOT を PATH
ROOTに追加する= Path(os.path.relpath(ROOT, Path.cwd())) # 相対

models.common から utils.datasets から DetectMultiBackend をインポートutils.generalから
IMG_FORMATS、VID_FORMATS、LoadImages、LoadStreams を                            インポート.plots import Annotator、colors、save_one_box from utils.torch_utils import select_device、time_sync




@torch.no_grad()
def run(weights=ROOT / 'yolov5s.pt', # model.pt パス
        vidpath='/home/ccf_disk/animal/test/', # file/dir/URL/glob, 0 for webcam
        data=ROOT / 'data/coco128.yaml', # dataset.yaml path
        imgsz=(640, 640), # 推論サイズ (高さ, 幅)
        conf_thres=0.6, # 信頼度しきい値
        iou_thres=0.45, # NMS IOU threshold
        max_det=1000, # 画像あたりの最大検出数
        device='', # cuda デバイス、つまり 0 または 0,1,2,3 または cpu
        view_img=False, # 結果を表示
        save_txt=False, # 結果を *.txt に保存
        save_conf =False, # --save-txt ラベルに信頼度を保存
        save_crop=False, # トリミングされた予測ボックスを保存する
        nosave=True, # 画像/動画を保存しない
        classes=None, # クラスでフィルタ: --class 0, または --class 0 2 3
        agnostic_nms=False, # クラスに依存しない NMS
        Augment=False, # 拡張推論
        Visualize=False, # 機能を視覚化
        update=False, # すべてのモデルを更新
        project='/home/ccf_disk/animal/video_animal', # 結果を project/name name
        ='test_1', # save に保存結果を project/name
        exist_ok=True, # 既存の project/name ok, インクリメントしない
        line_thickness=3, # バウンディング ボックスの太さ (ピクセル)
        hide_labels=False, # ラベルを隠す
        hide_conf=False, # 自信を隠す
        half=False, # FP16 半精度推論を使用
        dnn=False, # ONNX 推論に OpenCV DNN を使用
        ):
    vidpath = str(vidpath)
    ビデオ = os.listdir(vidpath)
    number = 0
    for video_nameビデオの場合:
        time1_start = time.time()
        so = vidpath + video_name
        number = number + 1
        print("第%d视频处理中" %number)
        source = str(so)
        save_c = 0
        keep = 0
        save_img = not nosave and not source.endswith('.txt') # 推論画像を保存
        is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
        is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
        webcam = source.isnumeric() または source. endwith('.txt') または (is_url で is_file ではない)
        if is_url and is_file:
            source = check_file(source) # download

        # ディレクトリ
        save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # インクリメント run
        (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir

        # モデルをロード
        device = select_device(device)
        model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data)
        stride, names, pt, jit, onnx, engine = model.stride, model.names, model.pt , model.jit, model.onnx, model.engine
        imgsz = check_img_size(imgsz, s=stride) # 画像サイズをチェック

        # Half half &= (pt または jit または onnx または engine) および device.type != 'cpu' #         pt または jit の場合、
        CUDA を使用する限定的なバックエンドでサポートされる FP16 :             model.model.half() if half else model.model.浮く()

        # Dataloader
        if webcam:
            view_img = check_imshow()
            cudnn.benchmark = True # 一定の画像サイズの推論を高速化するには True を設定します。
            dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt)
            bs = len(dataset) # batch_size
        else:
            dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
            bs = 1 # batch_size
        vid_path, vid_writer = [None] * bs, [None] * bs

        # 推論を実行
        model.warmup(imgsz=(1 if pt else bs, 3, *imgsz), half=half) # warmup
        dt, seen = [0.0, 0.0, 0.0], 0
        for path, im, im0s, vid_cap,データセットの s:
            flag = 0
            c = 1
            time1 = 6
            # t1 = time_sync()
            im = torch.from_numpy(im).to(device)
            im = im.half() if half else im.float() # uint8 to fp16/32
            im /= 255 # 0 - 255 から 0.0 - 1.0
            if len(im.shape) == 3:
                im = im[None] # バッチ ディム用に展開
            # t2 = time_sync()
            # dt[0] += t2 - t1

            # 推論

            視覚化 = increment_path(save_dir /パス(パス).stem
            ,             mkdir=True) if 視覚化
            else
t2

            # NMS
            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)

            # dt[2] += time_sync() - t3

            # 第 2 段階の分類器 (オプション)
            # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)


            # i の            予測を処理、det in enumerate(pred): #
                見た画像ごと += 1
                count = 0
                if webcam: # batch_size >= 1
                    p, im0, frame = path[i], im0s[i].copy() , dataset.count
                    s += f'{i}: '
                それ以外:
                    p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)

                p = Path(p) # to パス
                save_path = str(save_dir / p.name) # im.jpg
                txt_path = str(save_dir / 'labels' / p.stem) + (
                    '' if dataset.mode == 'image' else f'_{frame}') # im.txt
                s += '%gx%g ' % im.shape[2:] # 文字列を出力
                gn = torch.tensor(im0.shape)[[1, 0, 1                 , 0                 ] ] # 正規                     化
                ゲインim0サイズに



                    it[:, :4] = scale_coords(im.shape[2:], it[:, :4], im0.shape).round()

                    # c の結果を出力
                    in det[:, -1].unique():
                        n = (det[:, -1] == c).sum() # クラスごとの検出数
                        s += f"{n} {names [int(c)]}{'s' * (n > 1)}, " # 文字列に追加

                    #
                    *xyxy、conf、cls の結果を reversed(det) に書き込む:
                        count = 1
                        if save_txt: # ファイルに書き込む
                            xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view (-1).tolist() # 正規化された xywh
                            line = (cls, *xywh, conf) if save_conf else (cls, *xywh) #
                            open(txt_path + '.txt', 'a') as f:
                                f.write(('%g' * len(行)).rstrip() % 行 + '\n')

                        if save_img or save_crop or view_img: # bbox を画像に追加
                            c = int(cls) # integer class
                            label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}' )
                            annotator.box_label(xyxy, label, color=colors(c, True))
                            save_crop の場合:
                                save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg' , BGR=True)
                        box = xyxy        
                        c = int(cls) # 整数クラス
                        p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(ボックス[3]))
                        lw = max(round(sum(im0.shape) / 2 * 0.003), 2)
                        cv2.rectangle(im0, p1, p2, color=(0, 0, 255),
                                      太さ=max(round(sum(im0. shape) / 2 * 0.003), 2), lineType=cv2.LINE_AA)
                        label = (f'{names[c]} {conf:.2f}')
                        tf = max(lw - 1, 1)
                        w, h = cv2.getTextSize(label, 0, fontScale=lw / 3, thickness=tf)[0] # text width, height
                        outside = p1[1] - h - 3 >= 0 # ラベルはボックスの外側に収まります
                        cv2.putText(im0, label, (p1[0], p1[1] - 外側の場合は 2 p1[1] + h + 2), 0, lw / 3, (0, 0, 255)
                                    ,
                                    太さ=tf、線種=cv2.LINE_AA)

                # ストリーム結果
                im0 = annotator.result()
                if view_img:
                    cv2.imshow(str(p), im0)
                    cv2.waitKey(1) # 1 ミリ秒
                if (seen % time1 == 0):
                    if (count == 0) :
                        save_c = 0
                    else:
                        save_c = save_c + 1
                if(save_c>=4):
                    if keep == 0:
                        im0 = cv2.cvtColor(im0, cv2.COLOR_BGR2RGB)
                        frame = Image.fromarray(np.uint8(im0))
                        #print(save_path)
                        frame.save(str(save_path.split('.')[0]) + ".jpg")
                        keep = 1
                        shutil.copy(so, save_path)
                        print('have animal')
                        break
            else:
                continue
            break


            # # 結果を保存 (検出された画像)
            # if save_img:
            # if dataset.mode == 'image':
            # cv2.imwrite(save_path, im0)
            # else: # 'video' or 'stream'
            # if vid_path[i] != save_path: # 新しいビデオ
            # vid_path[i] = save_path
            # if isinstance(vid_writer[i], cv2.VideoWriter):
            # vid_writer[i].release() # 以前のビデオ ライターを解放
            # if vid_cap: # video
            # fps = vid_cap.get(cv2.CAP_PROP_FPS)
            # w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            # h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            # else: # ストリーム
            # fps, w, h = 30, im0.shape[1], im0.shape[0]
            # save_path = str(Path(save_path) .with_suffix('.mp4')) # 結果ビデオに *.mp4 サフィックスを強制
            # vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
            # vid_writer [i].write(im0)

            # 出力時間 (推論のみ)
            # LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')

        # 印刷結果
        # t = tuple(x / seen * 1E3 for x in dt) # 画像あたりの速度
        # LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
        save_txt または save_img の場合:
            s = f"\n{len(list(save_dir.glob('labels/*.txt')))} ラベルを { save_dir / 'labels'}" if save_txt else ''
            # LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
        if update:
            strip_optimizer(weights) # モデルの更新 (修正するため) SourceChangeWarning)

        time1_end = time.time()
        print('Video %d processing time' % number + str(time1_end-time1_start))
        # if bool == True:
        # shutil.copy(so, save_path)
        # else:
        # pass


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'weights/best.pt', help='model path( s)')
    parser.add_argument('--vidpath', type=str, default='/home/ccf_disk/animal/video/4-3/',
                        help='file/dir/URL/glob, ウェブカメラの場合は 0 ')
    parser.add_argument('--data', type=str, default=ROOT / 'data/myvoc.yaml', help='(オプション) dataset.yaml パス')
    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='推定サイズ h,w') parser.add_argument('--conf-
    thres '、タイプ = フロート、デフォルト = 0.75、ヘルプ ='信頼閾値')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
    parser.add_argument('--max-det', type=int, default=1000, help= '画像ごとの最大検出数')
    parser.add_argument('--device', default='', help='cuda デバイス、つまり 0 または 0,1,2,3 または cpu')
    parser.add_argument('--view -img', action='store_true', help='結果を表示')
    parser.add_argument('--save-txt', action='store_true', help='結果を *.txt に保存')
    parser.add_argument( '--save-conf', action='store_true', help='--save-txt ラベルに自信を保存')
    parser.add_argument('--save-crop', action='store_true', help='トリミングされた予測ボックスを保存')
    parser.add_argument('--nosave', action='store_true', help='画像/動画を保存しない')
    parser.add_argument('--classes', nargs='+', type=int, help='クラスによるフィルタ: --classes 0、または --classes 0 2 3') parser.add_argument('--
    agnostic- nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='拡張推論')
    parser.add_argument('--visualize' , action='store_true', help='visualize features')
    parser.add_argument('--update', action='store_true', help='すべてのモデルを更新')
    parser.add_argument('--project', default= '/home/ccf_disk/animal/video_animal_yolov5/', help='結果をプロジェクト/名前に保存')
    parser.add_argument('--name',default='4-3', help='結果をプロジェクト/名前に保存')
    parser.add_argument('--exist-ok', action='store_true', help='既存のプロジェクト/名前 OK, 増分しない')
    parser.add_argument('--line-thickness', default=3, type= int, help='境界ボックスの太さ (ピクセル)')
    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
    parser.add_argument('--hide -conf', default=False, action='store_true', help='hideconfidence')
    parser.add_argument('--half', action='store_true', help='FP16半精度推論を使用')
    パーサー. add_argument('--dnn', action='store_true', help='ONNX 推論に OpenCV DNN を使用')
    opt = パーサー。parse_args()
    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 #
    print_args(FILE.stem, opt)を展開
    リターンオプト


def main(opt):
    check_requirements(exclude=('tensorboard', 'thop'))
    run(**vars(opt))


if __name__ == "__main__":
    opt = parse_opt()
    メイン(opt)

ちょっと深いです)、ブログ投稿は初めてなので、浅めに記録します。

おすすめ

転載: blog.csdn.net/Xiashawuyanzu/article/details/126310868