torchvision 目標(biāo)檢測(cè)微調(diào)
本教程將使用Penn-Fudan Database for Pedestrian Detection and Segmentation 微調(diào) 預(yù)訓(xùn)練的Mask R-CNN 模型。 它包含 170 張圖片,345 個(gè)行人實(shí)例。
定義數(shù)據(jù)集
用于訓(xùn)練目標(biāo)檢測(cè)、實(shí)例分割和人物關(guān)鍵點(diǎn)檢測(cè)的參考腳本允許輕松支持添加新的自定義數(shù)據(jù)集。數(shù)據(jù)集應(yīng)繼承自標(biāo)準(zhǔn)的 torch.utils.data.dataset 類,并實(shí)現(xiàn) __len__ 和 __getitem__ 。
__getitem__ 需要返回:
image: PIL 圖像 (H, W)
target: 字典數(shù)據(jù),需要包含字段
boxes (FloatTensor[N, 4]): N 個(gè) Bounding box 的位置坐標(biāo) [x0, y0, x1, y1], 0~W, 0~H
labels (Int64Tensor[N]): 每個(gè) Bounding box 的類別標(biāo)簽,0 代表背景類。
image_id (Int64Tensor[1]): 圖像的標(biāo)簽 id,在數(shù)據(jù)集中是唯一的。
area (Tensor[N]): Bounding box 的面積,在 COCO 度量里使用,可以分別對(duì)不同大小的目標(biāo)進(jìn)行度量。
iscrowd (UInt8Tensor[N]): 如果 iscrowd=True 在評(píng)估時(shí)忽略。
(optionally) masks (UInt8Tensor[N, H, W]): 可選的 分割掩碼
(optionally) keypoints (FloatTensor[N, K, 3]): 對(duì)于 N 個(gè)目標(biāo)來(lái)說(shuō),包含 K 個(gè)關(guān)鍵點(diǎn) [x, y, visibility], visibility=0 表示關(guān)鍵點(diǎn)不可見。
如果模型可以返回上述方法,可以在訓(xùn)練、評(píng)估都能使用,可以用 pycocotools 里的腳本進(jìn)行評(píng)估。
pip install pycocotools 安裝工具。
關(guān)于 labels 有個(gè)說(shuō)明,模型默認(rèn) 0 為背景。如果數(shù)據(jù)集沒(méi)有背景類別,不需要在標(biāo)簽里添加 0 。 例如,假設(shè)有 cat 和 dog 兩類,定義了 1 表示 cat , 2 表示 dog , 如果一個(gè)圖像有兩個(gè)類別,類別的 tensor 為 [1, 2] 。
此外,如果希望在訓(xùn)練時(shí)使用縱橫比分組,那么建議實(shí)現(xiàn) get_height_and_width 方法,該方法將返回圖像的高度和寬度,如果未提供此方法,我們將通過(guò) __getitem__ 查詢數(shù)據(jù)集的所有元素,這會(huì)將圖像加載到內(nèi)存中,并且比提供自定義方法的速度慢。
為 PennFudan 寫自定義數(shù)據(jù)集
文件夾結(jié)構(gòu)如下:
PennFudanPed/
PedMasks/
FudanPed00001_mask.png
FudanPed00002_mask.png
FudanPed00003_mask.png
FudanPed00004_mask.png
...
PNGImages/
FudanPed00001.png
FudanPed00002.png
FudanPed00003.png
FudanPed00004.png
這是圖像的標(biāo)注信息,包含了?mask
?以及?bounding box
?。每個(gè)圖像都有對(duì)應(yīng)的分割掩碼,每個(gè)顏色代表不同的實(shí)例。
import os
import numpy as np
import torch
from PIL import Image
class PennFudanDataset(torch.utils.data.Dataset):
def __init__(self, root, transforms):
self.root = root
self.transforms = transforms
## 加載所有圖像,sort 保證他們能夠?qū)?yīng)起來(lái)
self.images = list(sorted(os.listdir(os.path.join(self.root, 'PNGImages'))))
self.masks = list(sorted(os.listdir(os.path.join(self.root, 'PedMasks'))))
def __getitem__(self, idx):
img_path = os.path.join(self.root, 'PNGImages', self.images[idx])
mask_path = os.path.join(self.root, 'PedMasks', self.masks[idx])
image = Image.open(img_path).convert('RGB')
## mask 圖像并沒(méi)有轉(zhuǎn)換為 RGB,里面存儲(chǔ)的是標(biāo)簽,0表示的是背景
mask = Image.open(mask_path)
# 轉(zhuǎn)換為 numpy
mask = np.array(mask)
# 實(shí)例解碼成不同的顏色
obj_ids = np.unique(mask)
# 移除背景
obj_ids = obj_ids[1:]
masks = mask == obj_ids[:, None, None]
# get bounding box coordinates for each mask
num_objs = len(obj_ids)
boxes = []
for i in range(num_objs):
pos = np.where(masks[i])
xmin = np.min(pos[1])
xmax = np.max(pos[1])
ymin = np.min(pos[0])
ymax = np.max(pos[0])
boxes.append([xmin, ymin, xmax, ymax])
# 轉(zhuǎn)換為 tensor
boxes = torch.as_tensor(boxes, dtype=torch.float32)
labels = torch.ones((num_objs,), dtype=torch.int64)
masks = torch.as_tensor(masks, dtype=torch.uint8)
image_id = torch.tensor([idx])
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
target = {}
target["boxes"] = boxes
target["labels"] = labels
target["masks"] = masks
target["image_id"] = image_id
target["area"] = area
target["iscrowd"] = iscrowd
if self.transforms is not None:
image, target = self.transforms(image, target)
return image, target
def __len__(self):
return len(self.images)
Lnton羚通專注于音視頻算法、算力、云平臺(tái)的高科技人工智能企業(yè)。 公司基于視頻分析技術(shù)、視頻智能傳輸技術(shù)、遠(yuǎn)程監(jiān)測(cè)技術(shù)以及智能語(yǔ)音融合技術(shù)等, 擁有多款可支持ONVIF、RTSP、GB/T28181等多協(xié)議、多路數(shù)的音視頻智能分析服務(wù)器/云平臺(tái)。文章來(lái)源:http://www.zghlxwxcb.cn/news/detail-658668.html
文章來(lái)源地址http://www.zghlxwxcb.cn/news/detail-658668.html
到了這里,關(guān)于Lnton羚通關(guān)于【PyTorch】教程:torchvision 目標(biāo)檢測(cè)微調(diào)的文章就介紹完了。如果您還想了解更多內(nèi)容,請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章,希望大家以后多多支持TOY模板網(wǎng)!