How to Run YOLOv7 Instance Segmentation Inference in Google Colab

我們使用 Google Colab 實操 YOLO 系列的最新版本 「YOLOv7」，輕鬆實現最新的實例分割模型。

如果對物件偵測和人體姿態估計有興趣的小夥伴，也可以點選下方連結唷！

How to Run YOLOv7 in Google Colab

How to Run YOLOv7 Human Post Estimation in Google Colab

在實際操作之前，我們先上一張巨石強森和傑森斯坦森的帥照吧！

二、Google Colab 設定

讓我們馬上開始使用 YOLOv7 吧！我們將使用 Google Colab 來實現。首先，配置 GPU 使其可以使用。“Runtime” -> “Change runtime type” -> 將“Hardware accelerator” 更改為 GPU。

接著安裝 Google Drive。

from google.colab import drive 
drive.mount('/content/drive')

創建新的資料夾。

!mkdir -p '/content/drive/My Drive/yolov7_segmentation/'
%cd '/content/drive/My Drive/yolov7_segmentation/'

從 YOLO 官方 git clone。

!git clone -b mask https://github.com/WongKinYiu/yolov7

然後導入所需的函式庫。這裡我們使用 pip 安裝。

%cd yolov7
!pip install -r requirements.txt

我們現在可以使用 YOLOv7 囉！開始出發吧！

從 Detectron2 官方 git clone。

!git clone https://github.com/facebookresearch/detectron2

更改資料夾名稱。

%mv ./detectron2 ./detectron

接下來，移動本次需要的 detectron2 子資料夾。

%mv ./detectron/detectron2 ./

安裝 Pytorch 。

%cd '/content/drive/My Drive/yolov7_segmentation/yolov7'
!pip install torchvision==0.10.1
!pip install torch==1.9.1

然後導入所需的函式庫。這裡我們使用 pip 安裝。

!pip install fvcore
!pip install omegaconf
!pip install fairscale
!pip install timm

準備工作現已完成。

使用範例圖片 inference

下載訓練好的模型。

!wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-mask.pt

首先，導入之前準備的每個函式庫。

import matplotlib.pyplot as plt
import torch
import cv2
import yaml
from torchvision import transforms
import numpy as np
from utils.datasets import letterbox
from utils.general import non_max_suppression_mask_conf
from detectron2.modeling.poolers import ROIPooler
from detectron2.structures import Boxes
from detectron2.utils.memory import retry_if_cuda_oom
from detectron2.layers import paste_masks_in_image

然後加載訓練好的模型。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
with open('data/hyp.scratch.mask.yaml') as f:
    hyp = yaml.load(f, Loader=yaml.FullLoader)
weigths = torch.load('yolov7-mask.pt')
model = weigths['model']
model = model.half().to(device)
_ = model.eval()

然後加載測試圖像。

image = cv2.imread('./inference/images/horses.jpg')  
image = letterbox(image, 960, stride=64, auto=True)[0]
image_ = image.copy()
image = transforms.ToTensor()(image)
image = torch.tensor(np.array([image.numpy()]))
image = image.to(device)
image = image.half()
output = model(image)

返回推理結果。

inf_out, train_out, attn, mask_iou, bases, sem_output = output['test'], output['bbox_and_cls'], output['attn'], output['mask_iou'], output['bases'], output['sem']
bases = torch.cat([bases, sem_output], dim=1)
nb, _, height, width = image.shape
names = model.names
pooler_scale = model.pooler_scale
pooler = ROIPooler(output_size=hyp['mask_resolution'], scales=(pooler_scale,), sampling_ratio=1, pooler_type='ROIAlignV2', canonical_level=2)
output, output_mask, output_mask_score, output_ac, output_ab = non_max_suppression_mask_conf(inf_out, attn, bases, pooler, hyp, conf_thres=0.25, iou_thres=0.65, merge=False, mask_iou=None)
pred, pred_masks = output[0], output_mask[0]
base = bases[0]
bboxes = Boxes(pred[:, :4])
original_pred_masks = pred_masks.view(-1, hyp['mask_resolution'], hyp['mask_resolution'])
pred_masks = retry_if_cuda_oom(paste_masks_in_image)( original_pred_masks, bboxes, (height, width), threshold=0.5)
pred_masks_np = pred_masks.detach().cpu().numpy()
pred_cls = pred[:, 5].detach().cpu().numpy()
pred_conf = pred[:, 4].detach().cpu().numpy()
nimg = image[0].permute(1, 2, 0) * 255
nimg = nimg.cpu().numpy().astype(np.uint8)
nimg = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)
nbboxes = bboxes.tensor.detach().cpu().numpy().astype(np.int)

根據最後獲得的結果對圖像進行遮罩。

pnimg = nimg.copy()
for one_mask, bbox, cls, conf in zip(pred_masks_np, nbboxes, pred_cls, pred_conf):
    if conf < 0.25:
        continue
    color = [np.random.randint(255), np.random.randint(255), np.random.randint(255)]              
    pnimg[one_mask] = pnimg[one_mask] * 0.5 + np.array(color, dtype=np.uint8) * 0.5

讓我們看看結果。

%matplotlib inline
plt.figure(figsize=(8,8))
plt.axis('off')
plt.imshow(pnimg)
plt.show()

添加物見偵測的訊息，同時顯示邊界框和分數。

pnimg = nimg.copy()
for one_mask, bbox, cls, conf in zip(pred_masks_np, nbboxes, pred_cls, pred_conf):
    if conf < 0.25:
        continue
    color = [np.random.randint(255), np.random.randint(255), np.random.randint(255)]
                        
                        
    pnimg[one_mask] = pnimg[one_mask] * 0.5 + np.array(color, dtype=np.uint8) * 0.5
    pnimg = cv2.rectangle(pnimg, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
    label = '%s %.3f' % (names[int(cls)], conf)
    t_size = cv2.getTextSize(label, 0, fontScale=0.5, thickness=1)[0]
    c2 = bbox[0] + t_size[0], bbox[1] - t_size[1] - 3
    pnimg = cv2.rectangle(pnimg, (bbox[0], bbox[1]), c2, color, -1, cv2.LINE_AA)  # filled
    pnimg = cv2.putText(pnimg, label, (bbox[0], bbox[1] - 2), 0, 0.5, [255, 255, 255], thickness=1, lineType=cv2.LINE_AA)
%matplotlib inline
plt.figure(figsize=(8,8))
plt.axis('off')
plt.imshow(pnimg)
plt.show()

使用任意影片 inference