【深度學習】【目標檢測】【Ultralytics-YOLO系列】YOLOV3核心文件detect.py解讀
文章目錄
- 【深度學習】【目標檢測】【Ultralytics-YOLO系列】YOLOV3核心文件detect.py解讀
- 前言
- if name == ‘main’
- parse_opt函數
- main函數
- run函數
- 不同命令參數的推理結果
- 常規推理命令
- 推理命令(新增save-txt參數)
- 推理命令(新增save-conf參數)
- 推理命令(新增save-crop參數)
- 推理命令(新增visualize參數)
- 總結
前言
在詳細解析YOLOV3網絡之前,首要任務是搭建Ultralytics–YOLOV3【Windows11下YOLOV3人臉檢測】所需的運行環境,并完成模型的訓練和測試,展開后續工作才有意義。
本博文對detect.py代碼進行解析,detect.py文件實現YOLOV3網絡模型的推理。其他代碼后續的博文將會陸續講解。這里只做YOLOV3相關模塊的代碼解析。
if name == ‘main’
Python腳本入口點:啟動程序,調用【detect.py】的parse_opt()函數負責解析命令行參數和【detect.py】的main()函數檢查依賴并調用推理函數。
if __name__ == "__main__":opt = parse_opt()main(opt)
parse_opt函數
解析命令行參數,用于配置YOLOv3模型的推理過程。
def parse_opt():"""解析命令行參數,用于配置YOLOv3模型的推理過程:return:命令行參數的對象"""parser = argparse.ArgumentParser()# weights:模型的權重地址,默認yolov3.pt nargs='+'參數選項用于指定命令行參數可以接受一個或多個值,并以列表的形式存儲parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov3.pt', help='model path(s)')# source:測試數據文件(圖片或視頻)的保存路徑,默認data/imagesparser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob, 0 for webcam')# imgsz:網絡輸入圖片的大小,默認640parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')# conf-thres:置信度閾值,默認0.25parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')# iou-thres:iou閾值(非極大值抑制NMS),默認0.45parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')# max-det:每張圖片最大的目標個數,默認1000parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')# device:執行的設備cuda(單卡0或者多卡0,1,2,3)或者cpuparser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')# view-img:是否展示預測之后的圖片或視頻,默認Falseparser.add_argument('--view-img', action='store_true', help='show results')# save-txt:是否保存預測的邊界框坐標到tx文件中,默認False,保存到runs/detect/expn/labels下每張圖片的txt文件parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')# save-conf:是否保存預測目標的置信度到tx文件中,默認False,保存到runs/detect/expn/labels下每張圖片的txt文件parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')# save-crop:是否需要從原圖中扣出預測到的目標剪切保存,在runs/detect/expn/crops下每個類別都有自己的文件夾保存對應的剪切圖片,默認Falseparser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')# nosave:是否不要保存預測后的圖片,默認False,保存預測后的圖片parser.add_argument('--nosave', action='store_true', help='do not save images/videos')# classes:nms是否是只保留特定的類default=[0,6,1,8,9,7],默認是None,保留所有類parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')# agnostic-nms:是否進行類別無關的nms,不區分類別直接對所有類別的邊界框進行統一處理,默認Falseparser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')# augment:是否使用數據增強進行推理,默認為Falseparser.add_argument('--augment', action='store_true', help='augmented inference')# visualize:是否可視化特征圖,默認為False,在runs/detect/expn下每張圖片都有自己的文件夾保存的不同階段的特征圖parser.add_argument('--visualize', action='store_true', help='visualize features')# update:對所有模型進行strip_optimizer操作,從保存模型文件中移除優化器狀態以及其他不必要的信息從而生成更輕量化的模型文件,默認為Falseparser.add_argument('--update', action='store_true', help='update all models')# project:當前測試結果放在哪個主文件夾下,默認runs/detectparser.add_argument('--project', default=ROOT / 'runs/detect', help='save results to project/name')# name:當前測試結果放在run/detect下的文件名,默認是exp,exp1,exp2.... 以此類推parser.add_argument('--name', default='exp', help='save results to project/name')# exist-ok:是否覆蓋已有結果,默認為Falseparser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')# line-thickness:畫邊界框的線條寬度,默認為3parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')# hide-labels:是否隱藏標簽信息,默認為Falseparser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')# hide-conf:是否隱藏置信度信息,默認為Falseparser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')# half:是否使用半精度Float16推理(縮短推理時間),默認是Falseparser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')# dnn:是否使 OpenCV DNN進行ONNX推理,默認為Falseparser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')# 解析命令行參數,并將結果存儲在opt對象中opt = parser.parse_args()# imgsz參數的長度為1則其值乘以2,否則保持不變opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1# 打印解析后的參數print_args(FILE.stem, opt) # FILE.stem是不含擴展名的文件名稱:這里是detect# 返回命令行參數的對象return opt
main函數
檢查環境和打印參數,并根據輸入參數啟動程序。
調用了【utils/general.py】的check_requirements函數;
調用了【detect.py】的run函數。
def main(opt):"""檢查環境和打印參數,并根據輸入參數啟動程序:param opt:命令行參數的對象:return:None"""# 檢查項目所需的依賴項:requrement.txt的包是否安裝check_requirements(exclude=('tensorboard', 'thop')) # 排除'tensorboard'和'thop'這兩個庫# 調用run啟動程序:命令行參數使用字典形式run(**vars(opt))
run函數
加載預訓練的YOLO模型,并對給定的圖像或視頻流執行目標檢測任務。
調用了【utils/general.py】的increment_path函數、check_img_size函數、check_imshow函數、increment_path函數、non_max_suppression函數、scale_coords函數、xyxy2xywh函數、colorstr函數和strip_optimizer函數;
調用了【utils/torch_utils.py】的select_device函數;
調用了【models/common.py】的DetectMultiBackend類;
調用了【utils/datasets.py】的LoadStreams類和LoadImages類;
調用了【utils/plots.py】的Annotator類和save_one_box函數。
- 數據準備: 主要負責數據輸入的準備和驗證,確保程序能夠正確處理不同來源的數據(如本地文件、網絡鏈接或攝像頭流)。
# ===================================== 1.數據準備 ===================================== source = str(source) # 輸入圖像/視頻的路徑 save_img = not nosave and not source.endswith('.txt') # 保存推理圖像標志:輸出只要不以.txt結尾且選擇保存預測結果,則都要保存預測后的圖片 # .suffix用于獲取文件路徑中的擴展名部分 is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) # 文件標志:檢查是否為圖像文件或視頻文件 is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://')) # URL標志:檢查是否為URL鏈接 # .isnumeric():是否是純數字字符串,數字通常用來表示攝像頭設備的索引 # .endswith('.txt') 是否以.txt結尾 # (is_url and not is_file)是否是URL且不是本地文件 webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file) # 檢查是否為攝像頭輸入 if is_url and is_file: # 網上的圖像文件或視頻文件source = check_file(source) # 下載文件 # ======================================================================================
- 預測結果保持路徑: 根據是否需要保存標簽文件來決定是否創建子目錄,它確保每次運行程序時,結果不會覆蓋之前的結果,而是保存到一個獨立的新目錄中。
# ====================================2.預測結果保持路徑======================================== save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # 遞增生成的路徑 (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # 創建路徑 # ===========================================================================================
- 模型加載: 加載目標檢測模型并配置推理參數,包括設備選擇、模型格式判斷、圖像尺寸調整和半精度設置等。
# ===================================== 3.模型加載 ===================================== device = select_device(device) # 選擇設備CUDA或CPU model = DetectMultiBackend(weights, device=device, dnn=dnn) # 加載模型 # stride:推理時所用到最大步長,默認為32;names:保存預測類別對象列表;pt:加載的是否是pytorch模型;jit:加載的是否是jit格式;onnx:加載的是否是onnx模型 stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx # 獲取模型屬性 # 確保輸入圖片的尺寸imgsz能整除stride=32 imgsz = check_img_size(imgsz, s=stride)# 半精度僅支持在CUDA上運行PyTorch模型 half &= pt and device.type != 'cpu' if pt:# 將模型參數從FP32轉換為FP16model.model.half() if half else model.model.float() # ====================================================================================
- 數據加載: 確保了程序能夠正確處理不同類型的輸入,無論是實時的攝像頭流還是靜態的圖像或視頻文件。
# ====================================4.數據加載======================================== # 不同的輸入源設置不同的加載方式 if webcam: # 使用攝像頭/網絡視頻流作為輸入view_img = check_imshow() # 檢測cv2.imshow()方法是否可以執行,不能執行則拋出異常cudnn.benchmark = True # 設置為True可以加速固定尺寸圖像的推理dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt and not jit) # 加載輸入數據流bs = len(dataset) # 批量大小,輸入視頻流的數量 else: # 獲取本地圖片/視頻作為輸入dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt and not jit) bs = 1 # 批量大小,單張圖片或視頻 # 用于保存輸出視頻路徑和文件 vid_path, vid_writer = [None] * bs, [None] * bs # ====================================================================================
- 模型推理: 模型推理與結果處理全流程,包含了從圖像預處理、模型前向傳播、NMS 后處理到最終的可視化和保存。
# ====================================5.模型推理======================================== # 模型預熱 if pt and device.type != 'cpu':# 對模型進行預熱,以提高后續推理速度(即用一個空張量運行一次前向傳播),目的是讓GPU提前分配內存優化計算圖等,從而提高后續推理速度# .to(device)將張量移動到指定設備(GPU或CPU);.type_as(...)匹配模型參數的數據類型# next(model.model.parameters())的作用是從模型的參數迭代器中獲取第一個參數model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.model.parameters()))) # 在后續的推理過程中會被用來記錄時間(三個階段的時間消耗)和處理的圖片或視頻幀數量 dt, seen = [0.0, 0.0, 0.0], 0 # 遍歷數據集中的每張圖片或視頻幀,并進行預處理 for path, im, im0s, vid_cap, s in dataset:# ----------------5.1數據預處理-------------------t1 = time_sync() # 高精度時間函數通常用于性能測試,作為預處理階段的起始時間im = torch.from_numpy(im).to(device) # 將NumPy數組轉換為PyTorch張量并將張量移動到指定設備im = im.half() if half else im.float() # 數據類型轉換,張量的數據類型使用FP16半精度或FP32單精度im /= 255 # 將像素值從[0,255]的范圍歸一化到[0.0,1.0]的范圍if len(im.shape) == 3: # 張量的形狀是三維im = im[None] # 增加一個批量維度(批量維度)t2 = time_sync() # 作為預處理階段的結束時間dt[0] += t2 - t1 # 記錄預處理階段總耗時# ----------------------------------------------# ----------------5.2執行推理-------------------# 生成一個用于保存可視化結果的路徑visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False# 輸入數據進行推理,返回預測結# augment:是否啟用測試時增強;visualize:是否保存網絡中間層的特征圖pred = model(im, augment=augment, visualize=visualize)t3 = time_sync() # 作為推理階段的結束時間dt[1] += t3 - t2 # 記錄推理階段總耗時# --------------------------------------------# ----------------5.3NMS-------------------# 對模型的預測結果進行非極大值抑制,去除冗余的檢測框,保留最優的結果pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det) # 非極大值抑制NMSdt[2] += time_sync() - t3 # 記錄后處理階段的總耗時# -----------------------------------------# Second-stage classifier (optional)# pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)# 遍歷每張圖片的預測結果for i, det in enumerate(pred): # 每次處理一張圖像或視頻幀seen += 1 # 統計總共處理的圖像/幀,自增1if webcam: # 攝像頭/網絡視頻流輸入 batch_size >= 1 多個攝像頭或網絡視頻流# p:圖像路徑;im0:未歸一化原始圖像;frame:當前幀編號p, im0, frame = path[i], im0s[i].copy(), dataset.count# 通常有多個視頻流輸入作為批量輸入,處理當前批次的那個圖片序號就表示處理的是那個視頻流的幀s += f'{i}: ' # 記錄日志信息else: # 本地圖像或視頻幀輸入# 從dataset對象中獲取名為frame的屬性值;沒有定義frame屬性,返回默認值0p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)p = Path(p) # 將路徑轉換為Path對象save_path = str(save_dir / p.name) # 保存檢測結果圖像的路徑# 視頻或攝像頭流,文件名則附加_幀號后綴txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # 保存檢測框標簽文件的路徑s += '%gx%g ' % im.shape[2:] # 日志字符串中添加當前圖像的尺寸信息# 原始圖像形狀為(H,W,C)或(H,W),[[1,0,1,0]]提取[W,H,W,H]gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # 將檢測框坐標從[0~1]歸一化格式轉換回原始圖像像素坐標的縮放因子imc = im0.copy() if save_crop else im0 # 保存裁剪目標,復制原始圖像,防止原圖被修改annotator = Annotator(im0, line_width=line_thickness, example=str(names))if len(det): # 有檢測結果則進行坐標映射# 將檢測框坐標從模型輸入尺寸640x640映射回原始圖像尺寸det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()# 統計檢測類別及數量# det[:, -1]:取出所有檢測目標的類別編號;.unique():返回張量中的唯一值,去重后的元素for c in det[:, -1].unique(): # 遍歷所有檢測到的類別編號n = (det[:, -1] == c).sum() # 統計每個類別出現次數s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # 構造日志字符串# *xyxy:左上角和右下角坐標[x1,y1,x2,y2];conf:置信度;cls:類別編號for *xyxy, conf, cls in reversed(det): # 用于遍歷檢測結果,通常優先處理高置信度的檢測框if save_txt: # 將檢測框信息寫入txt文件# xyxy2xywh()將[x1,y1,x2,y2]轉換為[xc,yc,w,h]格式xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # 將坐標歸一化到[0,1line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # 是否添加置信度# 寫入文件,每行一個檢測框with open(txt_path + '.txt', 'a') as f:f.write(('%g ' * len(line)).rstrip() % line + '\n')if save_img or save_crop or view_img: # 啟用繪圖功能,則在圖像上繪制檢測框和標簽c = int(cls) # 類別序號# hide_labels:不顯示標簽;hide_conf:不顯示置信度# 不顯示標簽則直接不顯示置信度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)) # 用于在圖像上繪制邊界框和標簽if save_crop: # 裁剪的目標區域save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True) # 保存裁剪的目標區域到指定路徑# 打印單幀圖像的目標檢測推理所用時間(不包括前處理和后處理)LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')im0 = annotator.result() # 獲取帶有邊界框標簽等信息的圖像if view_img: # 顯示圖像cv2.imshow(str(p), im0) # 將標注了檢測框的結果圖像顯示在窗口中,用于實時查看檢測效果cv2.waitKey(1) # 等待1毫秒,保持窗口更新,防止卡死if save_img: # 保存檢測結果圖像或視頻if dataset.mode == 'image': # 圖像模式cv2.imwrite(save_path, im0) # cv2.imwrite() 保存圖片else: # 視頻或流媒體模式if vid_path[i] != save_path: # 檢查當前視頻路徑是否與新的保存路徑不同vid_path[i] = save_path # 不同則更新路徑,開始處理新的視頻文件if isinstance(vid_writer[i], cv2.VideoWriter): # cv2.VideoWriter類型的對象vid_writer[i].release() # 釋放舊的資源,以避免沖突if vid_cap: # 視頻輸入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: # 攝像頭輸入# 設置默認幀率為30,寬度和高度為當前幀尺寸fps, w, h = 30, im0.shape[1], im0.shape[0]save_path += '.mp4' # .mp4格式保存# save_path:保存路徑;v2.VideoWriter_fourcc(*'mp4v'):MP4編碼格式# fps:幀率;(w, h):視頻分辨率vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) # 創建新的VideoWriter對象,用于寫入處理后的幀vid_writer[i].write(im0) # 將當前幀寫入視頻文件 # ====================================================================================
完整代碼
def run(weights=ROOT / 'yolov3.pt', # 訓練的權重路徑source=ROOT / 'data/images', # 圖像/視頻文件,圖像目錄,URL,0(攝像頭)imgsz=640, # 推理圖像分辨率conf_thres=0.25, # 置信度閾值iou_thres=0.45, # IOU閾值(非極大值抑制NMS)max_det=1000, # 每張圖像的最大檢測數device='', # CUDA單卡或多卡(0或0,1,2,3)或CPUview_img=False, # 顯示結果save_txt=False, # 將結果(類別和邊框大小位置)保存到txtsave_conf=False, # 將置信度保存到txtsave_crop=False, # 保存裁剪后的圖像nosave=False, # 不保存圖像/視頻結果classes=None, # 按類別過濾篩選保留agnostic_nms=False, # 類別無關的NMSaugment=False, # 增強推理visualize=False, # 可視化特征update=False, # 更新所有模型(簡潔化)project=ROOT / 'runs/detect', # 保存結果的路徑name='exp', # 保存結果的名稱exist_ok=False, # 覆蓋現有結果,不名稱遞增保存line_thickness=3, # 邊界框厚度(像素)hide_labels=False, # 輸出結果隱藏標簽hide_conf=False, # 輸出結果隱藏置信度half=False, # 使用FP16半精度推理dnn=False, # 使用OpenCV DNN進行ONNX推理):"""加載預訓練的YOLO模型,并對給定的圖像或視頻流執行目標檢測任務:param weights:訓練的權重路徑,可以使用自己訓練的權重,也可以使用官網提供的權重,默認官網的權重yolov3.pt:param source:測試數據,可以是圖片/視頻路徑,也可以是'0'(電腦自帶攝像頭),也可以是rtsp等視頻流,默認data/images:param imgsz: 網絡模型輸入圖片尺寸,默認的大小是640:param conf_thres:置信度閾值,默認為0.25:param iou_thres:nms的iou閾值,默認為0.45:param max_det:保留的最大檢測框數量,每張圖片中檢測目標的個數最多為1000:param device:設置設備CPU/CUDA/多CUDA,默認不設置:param view_img:是否界面展示檢測結果(圖片/視頻),默認False:param save_txt:是否將預測的框坐標以txt文件形式保存,默認False,使用時在路徑runs/detect/exp*/labels/*.txt下生成每張圖片預測的txt文件:param save_conf:是否將置信度conf也保存到txt中,默認False,使用時在路徑runs/detect/exp*/labels/*.txt下生成每張圖片預測的txt文件:param save_crop:是否保存裁剪預測框圖片,默認為False,使用時在runs/detect/exp*/crop/剪切類別文件夾/ 路徑下會保存每個接下來的目標:param nosave:不保存圖片/視頻, 默認保存因此不設置,在runs/detect/exp*/保存預測的結果:param classes:設置只保留某一部分類別[0,6,1,8,9,7],默認不設置,設置時在路徑runs/detect/exp*/下保存[0,6,1,8,9,7]對應的類別的圖片:param agnostic_nms:進行NMS去除不同類別之間的框,默認False:param augment: 測試時增強/多尺度預測:param visualize:是否可視化網絡層輸出特征:param update:為True則對所有模型進行strip_optimizer操作,去除pt文件中的優化器等信息,默認為False:param project:保存測試日志的文件夾路徑,默認runs/detect:param name:保存測試日志文件夾的名字,所以最終是保存在project/name中:param exist_ok:是否重新創建日志文件,False時重新創建文件:param line_thickness:畫框的線條粗細:param hide_labels:預測結果隱藏標簽:param hide_conf:預測結果隱藏置信度:param half:是否使用F16精度推理,半進度提高檢測速度:param dnn:用OpenCV DNN預測:return:None"""# ===================================== 1.數據準備 =====================================source = str(source) # 輸入圖像/視頻的路徑save_img = not nosave and not source.endswith('.txt') # 保存推理圖像標志:輸出只要不以.txt結尾且選擇保存預測結果,則都要保存預測后的圖片# .suffix用于獲取文件路徑中的擴展名部分is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) # 文件標志:檢查是否為圖像文件或視頻文件is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://')) # URL標志:檢查是否為URL鏈接# .isnumeric():是否是純數字字符串,數字通常用來表示攝像頭設備的索引# .endswith('.txt') 是否以.txt結尾# (is_url and not is_file)是否是URL且不是本地文件webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file) # 檢查是否為攝像頭輸入if is_url and is_file: # 網上的圖像文件或視頻文件source = check_file(source) # 下載文件# ======================================================================================# ====================================2.預測結果保持路徑========================================save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # 遞增生成的路徑(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # 創建路徑# ===========================================================================================# ===================================== 3.模型加載 =====================================device = select_device(device) # 選擇設備CUDA或CPUmodel = DetectMultiBackend(weights, device=device, dnn=dnn) # 加載模型# stride:推理時所用到最大步長,默認為32;names:保存預測類別對象列表;# pt:加載的是否是pytorch模型;jit:加載的是否是jit格式;onnx:加載的是否是onnx模型stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx # 獲取模型屬性# 確保輸入圖片的尺寸imgsz能整除stride=32imgsz = check_img_size(imgsz, s=stride)# 半精度僅支持在CUDA上運行PyTorch模型half &= pt and device.type != 'cpu'if pt:# 將模型參數從FP32轉換為FP16model.model.half() if half else model.model.float()# ====================================================================================# ====================================4.數據加載========================================# 不同的輸入源設置不同的加載方式if webcam: # 使用攝像頭/網絡視頻流作為輸入view_img = check_imshow() # 檢測cv2.imshow()方法是否可以執行,不能執行則拋出異常cudnn.benchmark = True # 設置為True可以加速固定尺寸圖像的推理dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt and not jit) # 加載輸入數據流bs = len(dataset) # 批量大小,輸入視頻流的數量else: # 獲取本地圖片/視頻作為輸入dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt and not jit)bs = 1 # 批量大小,單張圖片或視頻# 用于保存輸出視頻路徑和文件vid_path, vid_writer = [None] * bs, [None] * bs# ====================================================================================# ====================================5.模型推理========================================# 模型預熱if pt and device.type != 'cpu':# 對模型進行預熱,以提高后續推理速度(即用一個空張量運行一次前向傳播),目的是讓GPU提前分配內存優化計算圖等,從而提高后續推理速度# .to(device)將張量移動到指定設備(GPU或CPU);.type_as(...)匹配模型參數的數據類型# next(model.model.parameters())的作用是從模型的參數迭代器中獲取第一個參數model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.model.parameters())))# 在后續的推理過程中會被用來記錄時間(三個階段的時間消耗)和處理的圖片或視頻幀數量dt, seen = [0.0, 0.0, 0.0], 0# 遍歷數據集中的每張圖片或視頻幀,并進行預處理for path, im, im0s, vid_cap, s in dataset:# ----------------5.1數據預處理-------------------t1 = time_sync() # 高精度時間函數通常用于性能測試,作為預處理階段的起始時間im = torch.from_numpy(im).to(device) # 將NumPy數組轉換為PyTorch張量并將張量移動到指定設備im = im.half() if half else im.float() # 數據類型轉換,張量的數據類型使用FP16半精度或FP32單精度im /= 255 # 將像素值從[0,255]的范圍歸一化到[0.0,1.0]的范圍if len(im.shape) == 3: # 張量的形狀是三維im = im[None] # 增加一個批量維度(批量維度)t2 = time_sync() # 作為預處理階段的結束時間dt[0] += t2 - t1 # 記錄預處理階段總耗時# ----------------------------------------------# ----------------5.2執行推理-------------------# 生成一個用于保存可視化結果的路徑visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False# 輸入數據進行推理,返回預測結# augment:是否啟用測試時增強;visualize:是否保存網絡中間層的特征圖pred = model(im, augment=augment, visualize=visualize)t3 = time_sync() # 作為推理階段的結束時間dt[1] += t3 - t2 # 記錄推理階段總耗時# --------------------------------------------# ----------------5.3NMS-------------------# 對模型的預測結果進行非極大值抑制,去除冗余的檢測框,保留最優的結果pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det) # 非極大值抑制NMSdt[2] += time_sync() - t3 # 記錄后處理階段的總耗時# -----------------------------------------# Second-stage classifier (optional)# pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)# 遍歷每張圖片的預測結果for i, det in enumerate(pred): # 每次處理一張圖像或視頻幀seen += 1 # 統計總共處理的圖像/幀,自增1if webcam: # 攝像頭/網絡視頻流輸入 batch_size >= 1 多個攝像頭或網絡視頻流# p:圖像路徑;im0:未歸一化原始圖像;frame:當前幀編號p, im0, frame = path[i], im0s[i].copy(), dataset.count# 通常有多個視頻流輸入作為批量輸入,處理當前批次的那個圖片序號就表示處理的是那個視頻流的幀s += f'{i}: ' # 記錄日志信息else: # 本地圖像或視頻幀輸入# 從dataset對象中獲取名為frame的屬性值;沒有定義frame屬性,返回默認值0p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)p = Path(p) # 將路徑轉換為Path對象save_path = str(save_dir / p.name) # 保存檢測結果圖像的路徑# 視頻或攝像頭流,文件名則附加_幀號后綴txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # 保存檢測框標簽文件的路徑s += '%gx%g ' % im.shape[2:] # 日志字符串中添加當前圖像的尺寸信息# 原始圖像形狀為(H,W,C)或(H,W),[[1,0,1,0]]提取[W,H,W,H]gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # 將檢測框坐標從[0~1]歸一化格式轉換回原始圖像像素坐標的縮放因子imc = im0.copy() if save_crop else im0 # 保存裁剪目標,復制原始圖像,防止原圖被修改annotator = Annotator(im0, line_width=line_thickness, example=str(names))if len(det): # 有檢測結果則進行坐標映射# 將檢測框坐標從模型輸入尺寸640x640映射回原始圖像尺寸det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()# 統計檢測類別及數量# det[:, -1]:取出所有檢測目標的類別編號;.unique():返回張量中的唯一值,去重后的元素for c in det[:, -1].unique(): # 遍歷所有檢測到的類別編號n = (det[:, -1] == c).sum() # 統計每個類別出現次數s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # 構造日志字符串# *xyxy:左上角和右下角坐標[x1,y1,x2,y2];conf:置信度;cls:類別編號for *xyxy, conf, cls in reversed(det): # 用于遍歷檢測結果,通常優先處理高置信度的檢測框if save_txt: # 將檢測框信息寫入txt文件# xyxy2xywh()將[x1,y1,x2,y2]轉換為[xc,yc,w,h]格式xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # 將坐標歸一化到[0,1line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # 是否添加置信度# 寫入文件,每行一個檢測框with open(txt_path + '.txt', 'a') as f:f.write(('%g ' * len(line)).rstrip() % line + '\n')if save_img or save_crop or view_img: # 啟用繪圖功能,則在圖像上繪制檢測框和標簽c = int(cls) # 類別序號# hide_labels:不顯示標簽;hide_conf:不顯示置信度# 不顯示標簽則直接不顯示置信度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)) # 用于在圖像上繪制邊界框和標簽if save_crop: # 裁剪的目標區域save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True) # 保存裁剪的目標區域到指定路徑# 打印單幀圖像的目標檢測推理所用時間(不包括前處理和后處理)LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')im0 = annotator.result() # 獲取帶有邊界框標簽等信息的圖像if view_img: # 顯示圖像cv2.imshow(str(p), im0) # 將標注了檢測框的結果圖像顯示在窗口中,用于實時查看檢測效果cv2.waitKey(1) # 等待1毫秒,保持窗口更新,防止卡死if save_img: # 保存檢測結果圖像或視頻if dataset.mode == 'image': # 圖像模式cv2.imwrite(save_path, im0) # cv2.imwrite() 保存圖片else: # 視頻或流媒體模式if vid_path[i] != save_path: # 檢查當前視頻路徑是否與新的保存路徑不同vid_path[i] = save_path # 不同則更新路徑,開始處理新的視頻文件if isinstance(vid_writer[i], cv2.VideoWriter): # cv2.VideoWriter類型的對象vid_writer[i].release() # 釋放舊的資源,以避免沖突if vid_cap: # 視頻輸入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: # 攝像頭輸入# 設置默認幀率為30,寬度和高度為當前幀尺寸fps, w, h = 30, im0.shape[1], im0.shape[0]save_path += '.mp4' # .mp4格式保存# save_path:保存路徑;v2.VideoWriter_fourcc(*'mp4v'):MP4編碼格式# fps:幀率;(w, h):視頻分辨率vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) # 創建新的VideoWriter對象,用于寫入處理后的幀vid_writer[i].write(im0) # 將當前幀寫入視頻文件# ====================================================================================t = tuple(x / seen * 1E3 for x in dt) # 計算每個階段(預處理、推理、NMS)每張圖像所花費的平均時間# 使用日志記錄器輸出模型的速度信息LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)if save_txt or save_img: # 保存文本標簽或者保存圖像# save_dir.glob('labels/*.txt'):查找指定目錄下所有.txt文件;s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''# 顯示保存了多少個標簽文件,否則為空字符串LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")if update: # 設置了更新標志# 移除模型中的優化器狀態從而減小模型大小,并避免警告strip_optimizer(weights) # update model (to fix SourceChangeWarning)
不同命令參數的推理結果
一些常用命令參數的使用情況,并展示和說明其輸出結果。
常規推理命令
只進行最基礎的設置:
–weights指定模型權重文件路徑,用于加載訓練好的模型進行推理;
–source:指定要檢測的數據源,可以是圖片、視頻、攝像頭等;
–device:指定運行設備:0 表示使用 GPU(CUDA),也可以指定多個GPU如 0,1,或者使用 cpu;
–conf-thres:設置置信度閾值(confidence threshold),只有置信度高于該值的目標才會被保留;
–iou-thres:設置 NMS(非極大值抑制)的 IoU 閾值,用于去除重疊過多的預測框。
python detect.py --weights runs/train/exp/weights/best.pt --source data/images --device 0 --conf-thres 0.7 --iou-thres 0.3
控制臺輸出結果:
文件保存內容:
推理命令(新增save-txt參數)
–save-txt:保存預測邊界框信息到文本文件中。
python detect.py --weights runs/train/exp/weights/best.pt --source data/images --device 0 --conf-thres 0.7 --iou-thres 0.3 --save-txt
控制臺輸出結果:
文件保存內容:
0表示標簽序號;剩下的四個數值分別是框中心坐標以及框的尺寸。
推理命令(新增save-conf參數)
–save-conf:保存置信度信息到文本文件中。
python detect.py --weights runs/train/exp/weights/best.pt --source data/images --device 0 --conf-thres 0.7 --iou-thres 0.3 --save-txt --save-conf
使用–save-conf參數一定要先使用–save-txt參數,單獨使用不起作用,并且保存的數據是在同一個txt文件中。
控制臺輸出結果:
文件保存內容:
0表示標簽序號;中間四個數值分別是框中心坐標以及框的尺寸;最后一個數值是置信度。
推理命令(新增save-crop參數)
–save-crop:將檢測到的目標從原圖中裁剪出來并單獨保存。
python detect.py --weights runs/train/exp/weights/best.pt --source data/images --device 0 --conf-thres 0.7 --iou-thres 0.3 --save-txt --save-conf --save-crop
控制臺輸出結果:
文件保存內容:
推理命令(新增visualize參數)
–visualize:可視化模型中間層特征圖。
python detect.py --weights runs/train/exp/weights/best.pt --source data/images --device 0 --conf-thres 0.7 --iou-thres 0.3 --save-txt --save-conf --save-crop --visualize
控制臺輸出結果:
32/XXX:32是保存特征圖的數量,XXX是當前網絡層輸出特征圖的數量。
文件保存內容:
每個被檢測的圖像都單獨生成一個以圖像文件名命名的文件夾,保存不同網絡層的特征圖,這里每層限制保存32張特征圖。
總結
盡可能簡單、詳細的介紹了核心文件detect.py文件的作用:根據命令行參數設置YOLOv3模型的推理流程。