實(shí)戰(zhàn) | 基于YOLOv9+SAM實(shí)現(xiàn)動(dòng)態(tài)目標(biāo)檢測(cè)和分割(步驟 + 代碼)
共 13771字,需瀏覽 28分鐘
·
2024-04-18 10:05
點(diǎn)擊上方“小白學(xué)視覺”,選擇加"星標(biāo)"或“置頂”
重磅干貨,第一時(shí)間送達(dá)
導(dǎo) 讀
本文主要介紹基于YOLOv9+SAM實(shí)現(xiàn)動(dòng)態(tài)目標(biāo)檢測(cè)和分割,并給出詳細(xì)步驟和代碼。
背景介紹
在本文中,我們使用YOLOv9+SAM在RF100 Construction-Safety-2 數(shù)據(jù)集上實(shí)現(xiàn)自定義對(duì)象檢測(cè)模型。
這種集成不僅提高了在不同圖像中檢測(cè)和分割對(duì)象的準(zhǔn)確性和粒度,而且還擴(kuò)大了應(yīng)用范圍——從增強(qiáng)自動(dòng)駕駛系統(tǒng)到改進(jìn)醫(yī)學(xué)成像中的診斷過程。
通過利用 YOLOv9 的高效檢測(cè)功能和 SAM 以零樣本方式分割對(duì)象的能力,這種強(qiáng)大的組合最大限度地減少了對(duì)大量再訓(xùn)練或數(shù)據(jù)注釋的需求,使其成為一種多功能且可擴(kuò)展的解決方案。
YOLOv9簡(jiǎn)介
YOLOv9簡(jiǎn)介(You Only Look Once)
YOLOv9性能圖示
YOLOv9模型圖
SAM簡(jiǎn)介
SAM簡(jiǎn)介(Segment-Anything-Model)
SAM模型圖
關(guān)于數(shù)據(jù)集
該項(xiàng)目利用Roboflow 的RF100施工數(shù)據(jù)集,特別是Construction-Safety-2子集,來演示集成模型的功能。
https://blog.roboflow.com/roboflow-100/
實(shí)現(xiàn)步驟
實(shí)現(xiàn)步驟如下:
環(huán)境設(shè)置
下載 YOLOv9 和 SAM 的預(yù)訓(xùn)練模型權(quán)重
圖像推理
可視化和分析
獲取檢測(cè)結(jié)果
使用 SAM 進(jìn)行分割
環(huán)境設(shè)置
需要有 Google 帳戶才能訪問 Google Colab,這是一項(xiàng)免費(fèi)云服務(wù),為深度學(xué)習(xí)任務(wù)提供必要的計(jì)算資源,包括訪問高達(dá) 16 GB 的 GPU 和 TPU。
GPU狀態(tài)檢查
首先,我們確保 GPU 的可用性和就緒性,用于處理和運(yùn)行 YOLOv9+SAM 模型。
安裝 Google 云盤
接下來,如果您已經(jīng)下載了數(shù)據(jù)集,我們需要導(dǎo)航到存儲(chǔ)數(shù)據(jù)集的文件夾,否則我們可以直接使用 Roboflow 加載數(shù)據(jù)集。
from google.colab import drivedrive.mount('/content/drive')or%cd {HOME}/!pip install -q roboflowfrom roboflow import Roboflowrf = Roboflow(api_key="YOUR API KEY")project = rf.workspace("roboflow-100").project("construction-safety-gsnvb")dataset = project.version(2).download("yolov7")
配置 YOLOv9
數(shù)據(jù)集準(zhǔn)備好后,克隆 YOLOv9 存儲(chǔ)庫,然后切換到 YOLOv9 目錄并安裝所需的依賴項(xiàng),為對(duì)象檢測(cè)和分割任務(wù)做好準(zhǔn)備。
!git clone https://github.com/SkalskiP/yolov9.git%cd yolov9!pip3 install -r requirements.txt -q
顯示當(dāng)前目錄
將當(dāng)前工作目錄的路徑存儲(chǔ)在HOME變量中以供參考。
import osHOME = os.getcwd()print(HOME)
下載權(quán)重模型
讓我們?yōu)槟P蜋?quán)重創(chuàng)建一個(gè)目錄,并從 GitHub 上的發(fā)布頁面下載特定的 YOLOv9 和 GELAN 模型權(quán)重,這對(duì)于使用預(yù)訓(xùn)練參數(shù)初始化模型至關(guān)重要。
!mkdir -p {HOME}/weights!wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-c.pt!wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-e.pt!wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-c.pt!wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-e.pt
下載圖像進(jìn)行推理
為了使用 YOLOv9 權(quán)重進(jìn)行推理,我們必須設(shè)置一個(gè)數(shù)據(jù)目錄并下載一個(gè)示例圖像進(jìn)行處理,并在變量中設(shè)置該圖像的路徑SOURCE_IMAGE_PATH。
!mkdir -p {HOME}/data!wget -P {HOME}/data -q /content/drive/MyDrive/data/image9.jpegSOURCE_IMAGE_PATH = f"{HOME}/image9.jpeg"
使用自定義數(shù)據(jù)運(yùn)行檢測(cè)
之后,執(zhí)行detect.py指定參數(shù)對(duì)圖像進(jìn)行目標(biāo)檢測(cè),設(shè)置置信度閾值并保存檢測(cè)結(jié)果。這將創(chuàng)建一個(gè)包含 class_ids、邊界框坐標(biāo)和置信度分?jǐn)?shù)的文本文件,我們稍后將使用它。
!python detect.py --weights {HOME}/weights/gelan-c.pt --conf 0.1 --source /content/drive/MyDrive/data/image9.jpeg --device 0 --save-txt --save-conf
然后,我們利用 IPython 的顯示和圖像功能來展示指定路徑圖像中檢測(cè)到的對(duì)象,并進(jìn)行調(diào)整以獲得最佳觀看效果。
安裝 Ultralytics 包以訪問 YOLO 對(duì)象檢測(cè)模型實(shí)現(xiàn)和實(shí)用程序,不要忘記導(dǎo)入 YOLO 類以執(zhí)行對(duì)象檢測(cè)任務(wù)。
!pip install ultralyticsfrom ultralytics import YOLO
安裝Segment-Anything模型
現(xiàn)在讓我們安裝 Segment-Anything 庫并下載 SAM 模型的權(quán)重文件,為高質(zhì)量圖像分割任務(wù)做好準(zhǔn)備。
!pip install 'git+https://github.com/facebookresearch/segment-anything.git'!wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
提取檢測(cè)結(jié)果和置信度分?jǐn)?shù)
我們需要將 YOLOv9 檢測(cè)結(jié)果保存在上面的文本文件中來提取類 ID、置信度分?jǐn)?shù)和邊界框坐標(biāo)。這里的坐標(biāo)已經(jīng)標(biāo)準(zhǔn)化,所以我們先將它們轉(zhuǎn)換為圖像比例,然后打印它們進(jìn)行驗(yàn)證。
import cv2# Specify the path to your imageimage_path = '/content/drive/MyDrive/data/image9.jpeg'# Read the image to get its dimensionsimage = cv2.imread(image_path)image_width, _ = image.shapedetections_path = '/content/yolov9/runs/detect/exp/labels/image9.txt'bboxes = []class_ids = []conf_scores = []with open(detections_path, 'r') as file:for line in file:components = line.split()class_id = int(components[0])confidence = float(components[5])cy, w, h = [float(x) for x in components[1:5]]# Convert from normalized [0, 1] to image scalecx *= image_widthcy *= image_heightw *= image_widthh *= image_height# Convert the center x, y, width, and height to xmin, ymin, xmax, ymaxxmin = cx - w / 2ymin = cy - h / 2xmax = cx + w / 2ymax = cy + h / 2class_ids.append(class_id)ymin, xmax, ymax))conf_scores.append(confidence)# Display the resultsfor class_id, bbox, conf in zip(class_ids, bboxes, conf_scores):ID: {class_id}, Confidence: {conf:.2f}, BBox coordinates: {bbox}')
初始化 SAM 以進(jìn)行圖像分割
使用指定的預(yù)訓(xùn)練權(quán)重初始化 SAM 后,我們繼續(xù)從 SAM 模型注冊(cè)表中選擇模型類型以生成分段掩碼。
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictorsam_checkpoint = "/content/yolov9/sam_vit_h_4b8939.pth"model_type = "vit_h"sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)predictor = SamPredictor(sam)
加載圖像進(jìn)行分割
通過 OpenCV 庫,我們加載圖像以使用 SAM 進(jìn)行處理,為分割做好準(zhǔn)備。
import cv2image = cv2.cvtColor(cv2.imread('/content/drive/MyDrive/data/image9.jpeg'), cv2.COLOR_BGR2RGB)predictor.set_image(image)
結(jié)果可視化
為了可視化檢測(cè)和分割結(jié)果,我們必須使用 SAM 將邊界框轉(zhuǎn)換為分割掩模。我們隨機(jī)為類 ID 分配唯一的顏色,然后定義用于顯示掩碼、置信度分?jǐn)?shù)和邊界框的輔助函數(shù)。coco.yaml 文件用于將 class_ids 映射到類名。
import matplotlib.patches as patchesfrom matplotlib import pyplot as pltimport numpy as npimport yamlwith open('/content/yolov9/data/coco.yaml', 'r') as file:coco_data = yaml.safe_load(file)class_names = coco_data['names']for class_id, bbox, conf in zip(class_ids, bboxes, conf_scores):class_name = class_names[class_id]# print(f'Class ID: {class_id}, Class Name: {class_name}, BBox coordinates: {bbox}')color_map = {}for class_id in class_ids:color_map[class_id] = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)def show_mask(mask, ax, color):h, w = mask.shape[-2:]mask_image = mask.reshape(h, w, 1) * np.array(color).reshape(1, 1, -1)ax.imshow(mask_image)def show_box(box, label, conf_score, color, ax):x0, y0 = box[0], box[1]w, h = box[2] - box[0], box[3] - box[1]rect = plt.Rectangle((x0, y0), w, h, edgecolor=color, facecolor='none', lw=2)ax.add_patch(rect)label_offset = 10# Construct the label with the class name and confidence scorelabel_text = f'{label} {conf_score:.2f}'ax.text(x0, y0 - label_offset, label_text, color='black', fontsize=10, va='top', ha='left',bbox=dict(facecolor=color, alpha=0.7, edgecolor='none', boxstyle='square,pad=0.4'))plt.figure(figsize=(10, 10))ax = plt.gca()plt.imshow(image)# Display and process each bounding box with the corresponding maskfor class_id, bbox in zip(class_ids, bboxes):class_name = class_names[class_id]# print(f'Class ID: {class_id}, Class Name: {class_name}, BBox coordinates: {bbox}')color = color_map[class_id]input_box = np.array(bbox)# Generate the mask for the current bounding boxmasks, _, _ = predictor.predict(point_coords=None,point_labels=None,box=input_box,multimask_output=False,)show_mask(masks[0], ax, color=color)show_box(bbox, class_name, conf, color, ax)# Show the final plotplt.axis('off')plt.show()
提取掩碼
最后,我們生成一個(gè)合成圖像,在白色背景下突出顯示檢測(cè)到的對(duì)象,從分割蒙版創(chuàng)建聚合蒙版,并將其應(yīng)用到將原始圖像與白色背景混合以增強(qiáng)可視化。
aggregate_mask = np.zeros(image.shape[:2], dtype=np.uint8)# Generate and accumulate masks for all bounding boxesfor bbox in bboxes:input_box = np.array(bbox).reshape(1, 4)_, _ = predictor.predict(point_coords=None,point_labels=None,box=input_box,multimask_output=False,)aggregate_mask = np.where(masks[0] > 0.5, 1, aggregate_mask)# Convert the aggregate segmentation mask to a binary maskbinary_mask = np.where(aggregate_mask == 1, 1, 0)# Create a white background with the same size as the imagewhite_background = np.ones_like(image) * 255# Apply the binary mask to the original image# Where the binary mask is 0 (background), use white_background; otherwise, use the original imagenew_image = white_background * (1 - binary_mask[..., np.newaxis]) + image * binary_mask[..., np.newaxis]# Display the new image with the detections and white background=(10, 10))plt.imshow(new_image.astype(np.uint8))plt.axis('off')plt.show()
—THE END—
下載1:OpenCV-Contrib擴(kuò)展模塊中文版教程
在「小白學(xué)視覺」公眾號(hào)后臺(tái)回復(fù):擴(kuò)展模塊中文教程,即可下載全網(wǎng)第一份OpenCV擴(kuò)展模塊教程中文版,涵蓋擴(kuò)展模塊安裝、SFM算法、立體視覺、目標(biāo)跟蹤、生物視覺、超分辨率處理等二十多章內(nèi)容。
下載2:Python視覺實(shí)戰(zhàn)項(xiàng)目52講
在「小白學(xué)視覺」公眾號(hào)后臺(tái)回復(fù):Python視覺實(shí)戰(zhàn)項(xiàng)目,即可下載包括圖像分割、口罩檢測(cè)、車道線檢測(cè)、車輛計(jì)數(shù)、添加眼線、車牌識(shí)別、字符識(shí)別、情緒檢測(cè)、文本內(nèi)容提取、面部識(shí)別等31個(gè)視覺實(shí)戰(zhàn)項(xiàng)目,助力快速學(xué)校計(jì)算機(jī)視覺。
下載3:OpenCV實(shí)戰(zhàn)項(xiàng)目20講
在「小白學(xué)視覺」公眾號(hào)后臺(tái)回復(fù):OpenCV實(shí)戰(zhàn)項(xiàng)目20講,即可下載含有20個(gè)基于OpenCV實(shí)現(xiàn)20個(gè)實(shí)戰(zhàn)項(xiàng)目,實(shí)現(xiàn)OpenCV學(xué)習(xí)進(jìn)階。
交流群
歡迎加入公眾號(hào)讀者群一起和同行交流,目前有SLAM、三維視覺、傳感器、自動(dòng)駕駛、計(jì)算攝影、檢測(cè)、分割、識(shí)別、醫(yī)學(xué)影像、GAN、算法競(jìng)賽等微信群(以后會(huì)逐漸細(xì)分),請(qǐng)掃描下面微信號(hào)加群,備注:”昵稱+學(xué)校/公司+研究方向“,例如:”張三 + 上海交大 + 視覺SLAM“。請(qǐng)按照格式備注,否則不予通過。添加成功后會(huì)根據(jù)研究方向邀請(qǐng)進(jìn)入相關(guān)微信群。請(qǐng)勿在群內(nèi)發(fā)送廣告,否則會(huì)請(qǐng)出群,謝謝理解~
交流群
歡迎加入公眾號(hào)讀者群一起和同行交流,目前有SLAM、三維視覺、傳感器、自動(dòng)駕駛、計(jì)算攝影、檢測(cè)、分割、識(shí)別、醫(yī)學(xué)影像、GAN、算法競(jìng)賽等微信群(以后會(huì)逐漸細(xì)分),請(qǐng)掃描下面微信號(hào)加群,備注:”昵稱+學(xué)校/公司+研究方向“,例如:”張三 + 上海交大 + 視覺SLAM“。請(qǐng)按照格式備注,否則不予通過。添加成功后會(huì)根據(jù)研究方向邀請(qǐng)進(jìn)入相關(guān)微信群。請(qǐng)勿在群內(nèi)發(fā)送廣告,否則會(huì)請(qǐng)出群,謝謝理解~
