以下是yolov5算法中目标检测与跟踪模块的源代码：
```python
import torch
import numpy as np
from scipy.optimize import linear_sum_assignment
from utils.general import xywh2xyxy, non_max_suppression
class Tracker:
def __init__(self, opt):
self.opt = opt
self.nms_thres = opt.track_nms
self.max_time_lost = opt.track_max_lost
self.min_hits = opt.track_min_hits
self.track_buffer = []
self.id_count = 0
def update(self, results):
if results:
# get predicted locations from detections
pred_boxes = np.array([r['bbox'] for r in results])
scores = np.array([r['score'] for r in results])
classes = np.array([r['class'] for r in results])
areas = (pred_boxes[:, 2] - pred_boxes[:, 0]) * (pred_boxes[:, 3] - pred_boxes[:, 1])
# combine matched detections to assigned trackers
track_ids = np.zeros(len(self.track_buffer))
matched, unmatched_dets, unmatched_trks = self.associate_detections_to_trackers(
self.track_buffer, pred_boxes, self.nms_thres)
# update matched trackers with assigned detections
for m in matched:
track_ids[m[1]] = self.track_buffer[m[0]].update(
pred_boxes[m[1]], scores[m[1]], classes[m[1]], areas[m[1]])
# create and initialise new trackers for unmatched detections
for i in unmatched_dets:
track = Track(
pred_boxes[i], scores[i], classes[i], self.id_count, self.max_time_lost, areas[i])
self.track_buffer.append(track)
self.id_count += 1
# get trackers that are lost
lost_ids = [i for i, t in enumerate(self.track_buffer) if not t.is_confirmed()]
# if tracker is lost, check if it should be removed
del_tracks = []
for i in lost_ids:
if self.track_buffer[i].time_since_update > self.max_time_lost:
del_tracks.append(i)
self.track_buffer = [i for j, i in enumerate(self.track_buffer) if j not in del_tracks]
# get tracks that are confirmed
output_tracks = [t for t in self.track_buffer if t.is_confirmed()]
# get results
results = []
for track in output_tracks:
results.append({
'tracking_id': int(track.track_id),
'class': track.obj_class,
'bbox': [int(x) for x in track.to_tlbr()],
'score': track.score
})
return results
else:
# if no detections found, increase the time_since_update of all trackers and remove lost ones
del_tracks = []
for i, track in enumerate(self.track_buffer):
track.update([], -1, -1, -1)
if not track.is_confirmed():
del_tracks.append(i)
self.track_buffer = [i for j, i in enumerate(self.track_buffer) if j not in del_tracks]
return []
def associate_detections_to_trackers(self, trackers, detections, nms_thres):
# calculate iou between detections and trackers
iou_matrix = np.zeros((len(trackers), len(detections)), dtype=np.float32)
for i, trk in enumerate(trackers):
for j, det in enumerate(detections):
if trk.obj_class == det[4]:
bbox_trk = xywh2xyxy(np.array([trk.to_xywh()]))
bbox_det = xywh2xyxy(np.array([det[0:4]]))
iou_matrix[i, j] = bbox_iou(bbox_trk, bbox_det, CIoU=True)
# apply hungarian algorithm to find the best matches
matched_indices = linear_sum_assignment(-iou_matrix)
matched_indices = np.asarray(matched_indices)
matched_indices = np.transpose(matched_indices)
unmatched_trackers = []
for i, trk in enumerate(trackers):
if i not in matched_indices[:, 0]:
unmatched_trackers.append(i)
unmatched_detections = []
for i, detection in enumerate(detections):
if i not in matched_indices[:, 1]:
unmatched_detections.append(i)
# filter out matched with low IOU
matches = []
for m in matched_indices:
if iou_matrix[m[0], m[1]] < nms_thres:
unmatched_trackers.append(m[0])
unmatched_detections.append(m[1])
else:
matches.append(m.reshape(1, 2))
if len(matches) > 0:
matches = np.concatenate(matches, axis=0)
return matches, np.array(unmatched_detections), np.array(unmatched_trackers)
class Track:
def __init__(self, bbox, score, obj_class, track_id=None, max_lost=0, area=None):
self.bbox = bbox
self.track_id = track_id
self.max_lost = max_lost
self.time_since_update = 0
self.hits = 0
self.obj_class = obj_class
self.score = score
self.areas = [area]
self.prediction = np.asarray(bbox).reshape(1, 4)
if self.track_id is None:
self.track_id = Track._count
Track._count += 1
self.history = []
self.lost = False
self.age = 0
def update(self, bbox, score, obj_class, area):
self.time_since_update = 0
self.history = []
self.hits += 1
self.obj_class = obj_class
self.score = score
self.areas.append(area)
# exponential weighted averaging
alpha = 0.3
self.prediction = alpha * self.prediction + (1 - alpha) * np.asarray(bbox).reshape(1, 4)
return self.track_id
def predict(self):
self.prediction = self.prediction + np.array([0, 0, 0, 0])
self.time_since_update += 1
if self.time_since_update > self.max_lost:
self.lost = True
self.age += 1
if self.hits == 1:
self.history.append(self.bbox)
self.history.append(self.bbox)
else:
self.history.append(self.bbox)
return self.prediction[0]
def get_class(self):
return self.obj_class
def get_score(self):
return self.score
def get_area(self):
return self.areas[-1]
def to_tlbr(self):
return [int(self.bbox[0]), int(self.bbox[1]), int(self.bbox[2]), int(self.bbox[3])]
def to_xywh(self):
return [self.bbox[0], self.bbox[1], self.bbox[2] - self.bbox[0], self.bbox[3] - self.bbox[1]]
def is_confirmed(self):
return not self.lost
_count = 0
def bbox_iou(box1, box2, CIoU=False):
# get the coordinates of bounding boxes
b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]
# get the coordinates of the intersection rectangle
inter_rect_x1 = np.maximum(b1_x1, b2_x1)
inter_rect_y1 = np.maximum(b1_y1, b2_y1)
inter_rect_x2 = np.minimum(b1_x2, b2_x2)
inter_rect_y2 = np.minimum(b1_y2, b2_y2)
# intersection area
inter_area = np.clip(inter_rect_x2 - inter_rect_x1, 0, None) * np.clip(inter_rect_y2 - inter_rect_y1, 0, None)
# union area
b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
union_area = b1_area + b2_area - inter_area
# iou
iou = inter_area / union_area
if CIoU:
# convex w,h
c_w = np.maximum(b1_x2, b2_x2) - np.minimum(b1_x1, b2_x1)
c_h = np.maximum(b1_y2, b2_y2) - np.minimum(b1_y1, b2_y1)
# convex diagonal squared
c2 = c_w ** 2 + c_h ** 2
# centerpoints distance squared
rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4
# distance IoU
DIoU = iou - rho2 / c2
# Complete IoU Loss
v = (4 / (np.pi ** 2)) * np.power((np.arctan(b2_w / b2_h) - np.arctan(b1_w / b1_h)), 2)
alpha = v / (1 - iou + v)
CIoU_loss = DIoU - alpha * v
return CIoU_loss
return iou
```
其中，`Tracker`类主要实现了目标跟踪算法，`Track`类是对单个目标的封装，实现了目标的跟踪、预测、更新等操作，`bbox_iou`函数是计算两个目标框之间的IoU，用于目标关联。
扫一扫
5146
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
