切換語言為:簡體
快速用 Python 實現一個 AI 物體識別,五分鐘搞定!

快速用 Python 實現一個 AI 物體識別,五分鐘搞定!

  • 爱糖宝
  • 2024-11-03
  • 2048
  • 0
  • 0

一. 前言

近期在研究 AI 方面的應用 ,計算機視覺是其中的一大領域 ,應用場景也比較多。

這一篇主要嘗試弄清楚 : 一個 AI 視覺識別的 功能點 ,涉及到哪些知識 ?

二. 處理流程

2.1 前置知識點

在實現這個功能之前 ,先要明白幾個前置知識點 :

Python 層面需要的工具 :

元件工具 功能描述 作用
opencv-python 一個開源計算機視覺庫,提供影象和影片處理功能。 用於載入、處理和顯示影象,進行影象預處理,如縮放、裁剪和繪製。
tensorflow 一個開源機器學習框架,支援深度學習模型的構建和訓練。 用於構建和訓練物體識別模型,提供高層次的 API 和模型管理功能。
torch 一個開源深度學習框架,以動態計算圖著稱,適合研究和生產。 提供強大的張量計算功能,支援深度學習模型的訓練和推理。
torchvision PyTorch 的視覺工具包,提供常用的資料集、模型和影象變換。 用於載入標準資料集(如 COCO),並提供預訓練模型和影象預處理工具。
  • OpenCV 在前幾篇就有介紹 ,它作為計算機視覺庫 ,提供了很多圖形處理的功能。

  • 下面3個主要是 AI 領域的元件 ,用於使用 AI 的各項能力

模型和模型的選擇

我們在識別物體的過程中 ,圖片中的物體到底是什麼,這是需要一個 知識庫 來幫助給出答案的。 而模型就是那個知識庫。

  • 模型 是指一個數學表示或計算框架,旨在從輸入資料中學習模式和規律,並進行預測或分類

一個模型的生成包括 : 資料收集 , 資料預處理 ,資料建模 ,模型訓練 ,模型評估 ,模型測試 及 最終使用

關於模型的生成 ,不在本篇的討論中 ,展示也不會這一塊。 這一篇主要專注於使用。

  • 常用的物體識別模型

模型 作用說明 適合的場景
YOLO 實時物體檢測模型,透過全摺積網路實現快速檢測,適合大規模應用。 影片監控、自動駕駛、機器人視覺
SSD 單次檢測器,結合不同尺度的特徵圖進行物體檢測,速度快且準確。 移動裝置、實時影片分析
Faster R-CNN 採用區域提議網路(RPN)和深度摺積網路,實現高精度的物體檢測。 需要高精度的任務,如醫學影像分析
Mask R-CNN 在 Faster R-CNN 基礎上增加例項分割功能,能同時進行物體檢測和分割。 影象編輯、自動標註
RetinaNet 使用焦點損失(Focal Loss)處理類別不平衡問題,適合小物體檢測。 複雜背景中的物體檢測
EfficientDet 結合高效模型架構與多尺度特徵融合,資源有限的環境下表現出色。 嵌入式系統、移動裝置
CenterNet 基於中心點檢測的框架,簡單高效,適合實時檢測。 需要實時性的物體檢測應用

而本片主要使用的是 YOLO。

模型的使用

模型的使用有兩種方式 :

  • 一種是基於模型 + torch 的使用方式 , 也是最常規的用法

  • 一種是基於封裝的工具 ,實現快速的呼叫

2.2 模型的使用方式一

S1 : 下載 YoLo 模型

官方地址 :yolov5

// Python 外掛安裝
pip install torch torchvision torchaudio
pip install opencv-python
pip install matplotlib

// 克隆 YOLOv5 倉庫:
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

S2 : 編寫使用模型邏輯

import torch
import cv2
import matplotlib.pyplot as plt

# 載入 YOLOv5 模型(這裏使用預訓練的 'yolov5s' 模型)
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

# 讀取影象
image_path = 'C:\\Users\\Desktop\\test666.JPG'
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # OpenCV 使用 BGR 格式,轉換為 RGB

# 進行物體識別
results = model(image)

# 列印結果
results.print()  # 輸出識別資訊

# 視覺化結果
results.show()  # 顯示影象和檢測結果

# 儲存結果
results.save(save_dir='output')  # 儲存檢測結果到指定目錄

# 如果需要更詳細的處理
# 獲取檢測到的邊界框、標籤和置信度
boxes = results.xyxy[0]  # 結果格式: (x1, y1, x2, y2, confidence, class)
for box in boxes:
    x1, y1, x2, y2, conf, cls = box
    print(f"Detected {model.names[int(cls)]} with confidence {conf:.2f} at [{x1:.2f}, {y1:.2f}, {x2:.2f}, {y2:.2f}]")

# 還可以使用 plt 顯示影象
plt.imshow(image)
plt.axis('off')  # 不顯示座標軸
plt.show()

S3 : 測試

快速用 Python 實現一個 AI 物體識別,五分鐘搞定!

2.2 模型的使用方式二

第二種方式更簡單 ,例如 Ultralytics 就提供了一整套元件可以開箱即用。

S1 : Python 安裝 ultralytics

install ultralytics

S2 : 對物體進行識別

from ultralytics import YOLO
import cv2
import matplotlib.pyplot as plt

# 載入 YOLOv5 或 YOLOv8 模型
model = YOLO('yolov5s.pt')  # 或使用 'yolov8.pt'

# 讀取影象
image_path = 'C:\\Users\\Desktop\\test666.JPG'
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 轉換為 RGB 格式

# 進行物體識別
results = model.predict(source=image)

# 視覺化結果
# 使用 cv2 繪製邊界框
for result in results:
    boxes = result.boxes.xyxy  # 獲取邊界框資訊
    scores = result.boxes.conf  # 獲取置信度
    classes = result.boxes.cls  # 獲取類別索引

    for i in range(len(boxes)):
        box = boxes[i]
        score = scores[i]
        cls = classes[i]

        x1, y1, x2, y2 = box.numpy()  # 轉換為 numpy 陣列
        label = f"{model.names[int(cls)]} {score:.2f}"

        # 繪製邊界框
        cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 0), 2)
        cv2.putText(image, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

# 使用 matplotlib 顯示影象
plt.imshow(image)
plt.axis('off')  # 不顯示座標軸
plt.show()

  • 👉 執行 模型的時候 ,會自動從遠端下載模型

Downloading https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov5su.pt to 'yolov5su.pt'...

快速用 Python 實現一個 AI 物體識別,五分鐘搞定!

三. 擴充套件和發展

一般沒有特殊的要求 ,大部分情況下 ,我們會選擇框架直接使用。下面主要來介紹一下 ultralytics 的 API

載入模型 :

  • YOLO(model): 用於載入 YOLO 模型,可以是本地路徑或線上模型名稱。

  • predict(source): 對給定的影象、影片或資料夾進行物體檢測。

結果處理 :

  • results.print(): 列印檢測結果,包括類別和置信度。

  • results.xyxy: 返回檢測到的邊界框,格式為 (x1, y1, x2, y2, confidence, class)。

  • results.show(): 顯示帶有檢測框的影象。

  • results.save(save_dir): 將帶有檢測框的影象儲存到指定目錄。

其他的更復雜的 :

針對初始模型 ,還可以透過資料集和訓練來適配更復雜的場景 :

  • train(): 訓練模型,支援自定義資料集。

  • val(): 驗證模型效能,通常在訓練後進行。

  • Dataset(): 用於處理資料集,支援常見的資料格式。

  • create_dataloader(): 建立資料載入器,方便批次處理資料。

總結一下

基本上使用起來賊簡單 ,簡單照著文件操作一下就能實現。

但是 : 生成級使用不會這麼簡單 , 模型相當於一個通用的資料集 ,它只能適配常見的需求。


作者:志字輩小螞蟻
連結:https://juejin.cn/post/7432913408574177307

0則評論

您的電子郵件等資訊不會被公開,以下所有項目均必填

OK! You can skip this field.