使用PyTorch實(shí)現(xiàn)目標(biāo)檢測與跟蹤
點(diǎn)擊上方“小白學(xué)視覺”，選擇加"星標(biāo)"或“置頂”
重磅干貨，第一時(shí)間送達(dá)
引言
在昨天的文章中，我們介紹了如何在PyTorch中使用您自己的圖像來訓(xùn)練圖像分類器，然后使用它來進(jìn)行圖像識別。本文將展示如何使用預(yù)訓(xùn)練的分類器檢測圖像中的多個(gè)對象，并在視頻中跟蹤它們。
圖像中的目標(biāo)檢測
目標(biāo)檢測的算法有很多，YOLO跟SSD是現(xiàn)下最流行的算法。在本文中，我們將使用YOLOv3。在這里我們不會詳細(xì)討論YOLO，如果想對它有更多了解，可以參考下面的鏈接哦~（https://pjreddie.com/darknet/yolo/）
下面讓我們開始吧，依然從導(dǎo)入模塊開始：
from models import *from utils import *import os, sys, time, datetime, randomimport torchfrom torch.utils.data import DataLoaderfrom torchvision import datasets, transformsfrom torch.autograd import Variableimport matplotlib.pyplot as pltimport matplotlib.patches as patchesfrom PIL import Image
然后加載預(yù)訓(xùn)練的配置和權(quán)重，以及一些預(yù)定義的值，包括：圖像的尺寸、置信度閾值和非最大抑制閾值。
config_path='config/yolov3.cfg'weights_path='config/yolov3.weights'class_path='config/coco.names'img_size=416conf_thres=0.8nms_thres=0.4# Load model and weightsmodel = Darknet(config_path, img_size=img_size)model.load_weights(weights_path)model.cuda()model.eval()classes = utils.load_classes(class_path)Tensor = torch.cuda.FloatTensor
下面的函數(shù)將返回對指定圖像的檢測結(jié)果。
def detect_image(img):# scale and pad imageratio = min(img_size/img.size[0], img_size/img.size[1])imw = round(img.size[0] * ratio)imh = round(img.size[1] * ratio)img_transforms=transforms.Compose([transforms.Resize((imh,imw)),max(int((imh-imw)/2),0),(128,128,128)),transforms.ToTensor(),])# convert image to Tensorimage_tensor = img_transforms(img).float()image_tensor = image_tensor.unsqueeze_(0)input_img = Variable(image_tensor.type(Tensor))# run inference on the model and get detectionswith torch.no_grad():detections = model(input_img)detections = utils.non_max_suppression(detections, 80,nms_thres)return detections[0]
最后，讓我們通過加載一個(gè)圖像，獲取檢測結(jié)果，然后用檢測到的對象周圍的包圍框來顯示它。并為不同的類使用不同的顏色來區(qū)分。
# load image and get detectionsimg_path = "images/blueangels.jpg"prev_time = time.time()img = Image.open(img_path)detections = detect_image(img)inference_time = datetime.timedelta(seconds=time.time() - prev_time)print ('Inference Time: %s' % (inference_time))# Get bounding-box colorscmap = plt.get_cmap('tab20b')colors = [cmap(i) for i in np.linspace(0, 1, 20)]img = np.array(img)plt.figure()ax = plt.subplots(1, figsize=(12,9))ax.imshow(img)pad_x = max(img.shape[0] - img.shape[1], 0) * (img_size / max(img.shape))pad_y = max(img.shape[1] - img.shape[0], 0) * (img_size / max(img.shape))unpad_h = img_size - pad_yunpad_w = img_size - pad_xif detections is not None:unique_labels = detections[:, -1].cpu().unique()n_cls_preds = len(unique_labels)bbox_colors = random.sample(colors, n_cls_preds)# browse detections and draw bounding boxesfor x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:box_h = ((y2 - y1) / unpad_h) * img.shape[0]box_w = ((x2 - x1) / unpad_w) * img.shape[1]y1 = ((y1 - pad_y // 2) / unpad_h) * img.shape[0]x1 = ((x1 - pad_x // 2) / unpad_w) * img.shape[1]color = bbox_colors[int(np.where(unique_labels == int(cls_pred))[0])]bbox = patches.Rectangle((x1, y1), box_w, box_h,linewidth=2, edgecolor=color, facecolor='none')ax.add_patch(bbox)y1, s=classes[int(cls_pred)],color='white', verticalalignment='top',bbox={'color': color, 'pad': 0})plt.axis('off')# save image"-det.jpg"),bbox_inches='tight', pad_inches=0.0)plt.show()
下面是我們的一些檢測結(jié)果：
視頻中的目標(biāo)跟蹤
現(xiàn)在你知道了如何在圖像中檢測不同的物體。當(dāng)你在一個(gè)視頻中一幀一幀地看時(shí)，你會看到那些跟蹤框在移動。但是如果這些視頻幀中有多個(gè)對象，你如何知道一個(gè)幀中的對象是否與前一個(gè)幀中的對象相同？這被稱為目標(biāo)跟蹤，它使用多次檢測來識別一個(gè)特定的對象。
有多種算法可以做到這一點(diǎn)，在本文中決定使用SORT(Simple Online and Realtime Tracking)，它使用Kalman濾波器預(yù)測先前識別的目標(biāo)的軌跡，并將其與新的檢測結(jié)果進(jìn)行匹配，非常方便且速度很快。
現(xiàn)在開始編寫代碼，前3個(gè)代碼段將與單幅圖像檢測中的代碼段相同，因?yàn)樗鼈兲幚淼氖窃趩螏汐@得 YOLO 檢測。差異在最后一部分出現(xiàn)，對于每個(gè)檢測，我們調(diào)用 Sort 對象的 Update 函數(shù)，以獲得對圖像中對象的引用。因此，與前面示例中的常規(guī)檢測(包括邊界框的坐標(biāo)和類預(yù)測)不同，我們將獲得跟蹤的對象，除了上面的參數(shù)，還包括一個(gè)對象 ID。并且需要使用OpenCV來讀取視頻并顯示視頻幀。
videopath = 'video/intersection.mp4'inlineimport cv2from IPython.display import clear_outputcmap = plt.get_cmap('tab20b')colors = [cmap(i)[:3] for i in np.linspace(0, 1, 20)]# initialize Sort object and video capturefrom sort import *vid = cv2.VideoCapture(videopath)mot_tracker = Sort()#while(True):for ii in range(40):frame = vid.read()frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)pilimg = Image.fromarray(frame)detections = detect_image(pilimg)img = np.array(pilimg)pad_x = max(img.shape[0] - img.shape[1], 0) */ max(img.shape))pad_y = max(img.shape[1] - img.shape[0], 0) */ max(img.shape))unpad_h = img_size - pad_yunpad_w = img_size - pad_xif detections is not None:tracked_objects = mot_tracker.update(detections.cpu())unique_labels = detections[:, -1].cpu().unique()n_cls_preds = len(unique_labels)for x1, y1, x2, y2, obj_id, cls_pred in tracked_objects:box_h = int(((y2 - y1) / unpad_h) * img.shape[0])box_w = int(((x2 - x1) / unpad_w) * img.shape[1])y1 = int(((y1 - pad_y // 2) / unpad_h) * img.shape[0])x1 = int(((x1 - pad_x // 2) / unpad_w) * img.shape[1])color = colors[int(obj_id) % len(colors)]color = [i * 255 for i in color]cls = classes[int(cls_pred)](x1, y1), (x1+box_w, y1+box_h),4)(x1, y1-35), (x1+len(cls)*19+60,color, -1)cls + "-" + str(int(obj_id)),y1 - 10), cv2.FONT_HERSHEY_SIMPLEX,(255,255,255), 3)fig=figure(figsize=(12, 8))Stream")imshow(frame)show()=True)
下面讓我們來看一下處理的結(jié)果：