使用PyTorch实现目标检测与跟踪！这不有趣多了！

引言

在昨天的文章中，我们介绍了如何在PyTorch中使用您自己的图像来训练图像分类器，然后使用它来进行图像识别。本文将展示如何使用预训练的分类器检测图像中的多个对象，并在视频中跟踪它们。

图像中的目标检测

目标检测的算法有很多，YOLO跟SSD是现下最流行的算法。在本文中，我们将使用YOLOv3。在这里我们不会详细讨论YOLO，如果想对它有更多了解，可以参考下面的链接哦~（https://pjreddie.com/darknet/yolo/）

下面让我们开始吧，依然从导入模块开始：

 

1. from models import *

2. from utils import *

3. import os, sys, time, datetime, random

4. import torch

5. from torch.utils.data import DataLoader

6. from torchvision import datasets, transforms

7. from torch.autograd import Variable

8. import matplotlib.pyplot as plt

9. import matplotlib.patches as patches

10. from PIL import Image

 

然后加载预训练的配置和权重，以及一些预定义的值，包括：图像的尺寸、置信度阈值和非最大抑制阈值。

 

1. config_path='config/yolov3.cfg'

2. weights_path='config/yolov3.weights'

3. class_path='config/coco.names'

4. img_size=416

5. conf_thres=0.8

6. nms_thres=0.4

7. # Load model and weights

8. model = Darknet(config_path, img_size=img_size)

9. model.load_weights(weights_path)

10. model.cuda()

11. model.eval()

12. classes = utils.load_classes(class_path)

13. Tensor = torch.cuda.FloatTensor

 

下面的函数将返回对指定图像的检测结果。

 

1. def detect_image(img):

2. # scale and pad image

3. ratio = min(img_size/img.size[0], img_size/img.size[1])

4. imw = round(img.size[0] * ratio)

5. imh = round(img.size[1] * ratio)

6. img_transforms=transforms.Compose([transforms.Resize((imh,imw)),

7. transforms.Pad((max(int((imh-imw)/2),0),

8. max(int((imw-imh)/2),0), max(int((imh-imw)/2),0),

9. max(int((imw-imh)/2),0)), (128,128,128)),

10. transforms.ToTensor(),

11. ])

12. # convert image to Tensor

13. image_tensor = img_transforms(img).float()

14. image_tensor = image_tensor.unsqueeze_(0)

15. input_img = Variable(image_tensor.type(Tensor))

16. # run inference on the model and get detections

17. with torch.no_grad():

18. detections = model(input_img)

19. detections = utils.non_max_suppression(detections, 80,

20. conf_thres, nms_thres)

21. return detections[0]

 

最后，让我们通过加载一个图像，获取检测结果，然后用检测到的对象周围的包围框来显示它。并为不同的类使用不同的颜色来区分。

 

1. # load image and get detections

2. img_path = "images/blueangels.jpg"

3. prev_time = time.time()

4. img = Image.open(img_path)

5. detections = detect_image(img)

6. inference_time = datetime.timedelta(seconds=time.time() - prev_time)

7. print ('Inference Time: %s' % (inference_time))

8. # Get bounding-box colors

9. cmap = plt.get_cmap('tab20b')

10. colors = [cmap(i) for i in np.linspace(0, 1, 20)]

11. img = np.array(img)

12. plt.figure()

13. fig, ax = plt.subplots(1, figsize=(12,9))

14. ax.imshow(img)

15. pad_x = max(img.shape[0] - img.shape[1], 0) * (img_size / max(img.shape))

16. pad_y = max(img.shape[1] - img.shape[0], 0) * (img_size / max(img.shape))

17. unpad_h = img_size - pad_y

18. unpad_w = img_size - pad_x

19. if detections is not None:

20. unique_labels = detections[:, -1].cpu().unique()

21. n_cls_preds = len(unique_labels)

22. bbox_colors = random.sample(colors, n_cls_preds)

23. # browse detections and draw bounding boxes

24. for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:

25. box_h = ((y2 - y1) / unpad_h) * img.shape[0]

26. box_w = ((x2 - x1) / unpad_w) * img.shape[1]

27. y1 = ((y1 - pad_y // 2) / unpad_h) * img.shape[0]

28. x1 = ((x1 - pad_x // 2) / unpad_w) * img.shape[1]

29. color = bbox_colors[int(np.where(

30. unique_labels == int(cls_pred))[0])]

31. bbox = patches.Rectangle((x1, y1), box_w, box_h,

32. linewidth=2, edgecolor=color, facecolor='none')

33. ax.add_patch(bbox)

34. plt.text(x1, y1, s=classes[int(cls_pred)],

35. color='white', verticalalignment='top',

36. bbox={'color': color, 'pad': 0})

37. plt.axis('off')

38. # save image

39. plt.savefig(img_path.replace(".jpg", "-det.jpg"),

40. bbox_inches='tight', pad_inches=0.0)

41. plt.show()

 

下面是我们的一些检测结果：

 

 

 

 

 

 

视频中的目标跟踪

现在你知道了如何在图像中检测不同的物体。当你在一个视频中一帧一帧地看时，你会看到那些跟踪框在移动。但是如果这些视频帧中有多个对象，你如何知道一个帧中的对象是否与前一个帧中的对象相同？这被称为目标跟踪，它使用多次检测来识别一个特定的对象。

有多种算法可以做到这一点，在本文中决定使用SORT(Simple Online and Realtime Tracking)，它使用Kalman滤波器预测先前识别的目标的轨迹，并将其与新的检测结果进行匹配，非常方便且速度很快。

现在开始编写代码，前3个代码段将与单幅图像检测中的代码段相同，因为它们处理的是在单帧上获得 YOLO 检测。差异在最后一部分出现，对于每个检测，我们调用 Sort 对象的 Update 函数，以获得对图像中对象的引用。因此，与前面示例中的常规检测(包括边界框的坐标和类预测)不同，我们将获得跟踪的对象，除了上面的参数，还包括一个对象 ID。并且需要使用OpenCV来读取视频并显示视频帧。

 

1. videopath = 'video/interp.mp4'

2. %pylab inline

3. import cv2

4. from IPython.display import clear_output

5. cmap = plt.get_cmap('tab20b')

6. colors = [cmap(i)[:3] for i in np.linspace(0, 1, 20)]

7. # initialize Sort object and video capture

8. from sort import *

9. vid = cv2.VideoCapture(videopath)

10. mot_tracker = Sort()

11. #while(True):

12. for ii in range(40):

13. ret, frame = vid.read()

14. frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

15. pilimg = Image.fromarray(frame)

16. detections = detect_image(pilimg)

17. img = np.array(pilimg)

18. pad_x = max(img.shape[0] - img.shape[1], 0) *

19. (img_size / max(img.shape))

20. pad_y = max(img.shape[1] - img.shape[0], 0) *

21. (img_size / max(img.shape))

22. unpad_h = img_size - pad_y

23. unpad_w = img_size - pad_x

24. if detections is not None:

25. tracked_objects = mot_tracker.update(detections.cpu())

26. unique_labels = detections[:, -1].cpu().unique()

27. n_cls_preds = len(unique_labels)

28. for x1, y1, x2, y2, obj_id, cls_pred in tracked_objects:

29. box_h = int(((y2 - y1) / unpad_h) * img.shape[0])

30. box_w = int(((x2 - x1) / unpad_w) * img.shape[1])

31. y1 = int(((y1 - pad_y // 2) / unpad_h) * img.shape[0])

32. x1 = int(((x1 - pad_x // 2) / unpad_w) * img.shape[1])

33. color = colors[int(obj_id) % len(colors)]

34. color = [i * 255 for i in color]

35. cls = classes[int(cls_pred)]

36. cv2.rectangle(frame, (x1, y1), (x1+box_w, y1+box_h),

37. color, 4)

38. cv2.rectangle(frame, (x1, y1-35), (x1+len(cls)*19+60,

39. y1), color, -1)

40. cv2.putText(frame, cls + "-" + str(int(obj_id)),

41. (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX,

42. 1, (255,255,255), 3)

43. fig=figure(figsize=(12, 8))

44. title("Video Stream")

45. imshow(frame)

46. show()

47. clear_output(wait=True)

PS：如有需要Python学习资料的小伙伴可以加点击下方链接自行获取

python免费学习资料以及群交流解答点击即可加入

下面让我们来看一下处理的结果：