def train(path, img_size, cfg='yolov5s.yaml', bs=2, one_batch_training=False):
device = 'cuda' if torch.cuda.is_available() else 'cpu'
dls, dsets = create_dataloaders(path, img_size, bs, device,
one_batch_training)
n_classes = len(dls.vocab)
model = Model(cfg=check_file(cfg), nc=n_classes)
if 'cuda' == device:
model.cuda()
hyp['cls'] *= n_classes / 80. # scale coco-tuned hyp['cls'] to current dataset
model.nc = n_classes # attach number of classes to model
model.hyp = hyp
model.gr = 1.0
learner = Learner(dls,
model,
loss_func=partial(compute_loss, model=model),
cbs=[EvaluatorCallback()])
with learner.no_bar():
learner.fit_one_cycle(args.epochs, lr_max=3e-3)
learner.save('/content/model_temp')
learner.export(fname='/content/learner_05_02_2021.pkl')
return learner
class JDETracker(object):
def __init__(self, opt, frame_rate=30):
self.opt = opt
self.model = Model(opt.cfg, nc=1)
ckpt = torch.load(opt.weights, map_location='cpu')
ckpt['model'] = {
k: v
for k, v in ckpt['model'].float().state_dict().items()
if k in self.model.state_dict()
and self.model.state_dict()[k].shape == v.shape
}
self.model.load_state_dict(ckpt['model'], strict=False)
self.model.cuda().half().eval()
self.tracked_stracks = [] # type: list[STrack]
self.lost_stracks = [] # type: list[STrack]
self.removed_stracks = [] # type: list[STrack]
self.frame_id = 0
self.det_thresh = opt.conf_thres
self.buffer_size = int(frame_rate / 30.0 * opt.track_buffer)
self.max_time_lost = self.buffer_size
self.kalman_filter = KalmanFilter()
def update(self, im_blob, img0, path):
self.frame_id += 1
activated_starcks = [
] # for storing active tracks, for the current frame
refind_stracks = [
] # Lost Tracks whose detections are obtained in the current frame
lost_stracks = [
] # The tracks which are not obtained in the current frame but are not removed.(Lost for some time lesser than the threshold for removing)
removed_stracks = []
t1 = time.time()
''' Step 1: Network forward, get detections & embeddings'''
with torch.no_grad():
pred = self.model(im_blob)[0]
# pred is tensor of all the proposals (default number of proposals: 54264). Proposals have information associated with the bounding box and embeddings
pred = pred[pred[:, :, 4] > self.opt.conf_thres]
# pred now has lesser number of proposals. Proposals rejected on basis of object confidence score
if len(pred) > 0:
dets = non_max_suppression(pred.unsqueeze(0), self.opt.conf_thres,
self.opt.nms_thres)[0].cpu()
# Final proposals are obtained in dets. Information of bounding box and embeddings also included
# Next step changes the detection scales
scale_coords(self.opt.img_size, dets[:, :4], img0.shape).round()
'''Detections is list of (x1, y1, x2, y2, object_conf, class_score, class_pred)'''
# class_pred is the embeddings.
detections = [
STrack(STrack.tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f.numpy(), 30)
for (tlbrs, f) in zip(dets[:, :5], dets[:, 6:])
]
else:
detections = []
t2 = time.time()
# print('Forward: {} s'.format(t2-t1))
''' Add newly detected tracklets to tracked_stracks'''
unconfirmed = []
tracked_stracks = [] # type: list[STrack]
for track in self.tracked_stracks:
if not track.is_activated:
# previous tracks which are not active in the current frame are added in unconfirmed list
unconfirmed.append(track)
# print("Should not be here, in unconfirmed")
else:
# Active tracks are added to the local list 'tracked_stracks'
tracked_stracks.append(track)
''' Step 2: First association, with embedding'''
# Combining currently tracked_stracks and lost_stracks
strack_pool = joint_stracks(tracked_stracks, self.lost_stracks)
# Predict the current location with KF
STrack.multi_predict(strack_pool, self.kalman_filter)
dists = matching.embedding_distance(strack_pool, detections)
# dists = matching.gate_cost_matrix(self.kalman_filter, dists, strack_pool, detections)
dists = matching.fuse_motion(self.kalman_filter, dists, strack_pool,
detections)
# The dists is the list of distances of the detection with the tracks in strack_pool
matches, u_track, u_detection = matching.linear_assignment(dists,
thresh=0.7)
# The matches is the array for corresponding matches of the detection with the corresponding strack_pool
for itracked, idet in matches:
# itracked is the id of the track and idet is the detection
track = strack_pool[itracked]
det = detections[idet]
if track.state == TrackState.Tracked:
# If the track is active, add the detection to the track
track.update(detections[idet], self.frame_id)
activated_starcks.append(track)
else:
# We have obtained a detection from a track which is not active, hence put the track in refind_stracks list
track.re_activate(det, self.frame_id, new_id=False)
refind_stracks.append(track)
# None of the steps below happen if there are no undetected tracks.
''' Step 3: Second association, with IOU'''
detections = [detections[i] for i in u_detection]
# detections is now a list of the unmatched detections
r_tracked_stracks = [
] # This is container for stracks which were tracked till the
# previous frame but no detection was found for it in the current frame
for i in u_track:
if strack_pool[i].state == TrackState.Tracked:
r_tracked_stracks.append(strack_pool[i])
dists = matching.iou_distance(r_tracked_stracks, detections)
matches, u_track, u_detection = matching.linear_assignment(dists,
thresh=0.5)
# matches is the list of detections which matched with corresponding tracks by IOU distance method
for itracked, idet in matches:
track = r_tracked_stracks[itracked]
det = detections[idet]
if track.state == TrackState.Tracked:
track.update(det, self.frame_id)
activated_starcks.append(track)
else:
track.re_activate(det, self.frame_id, new_id=False)
refind_stracks.append(track)
# Same process done for some unmatched detections, but now considering IOU_distance as measure
for it in u_track:
track = r_tracked_stracks[it]
if not track.state == TrackState.Lost:
track.mark_lost()
lost_stracks.append(track)
# If no detections are obtained for tracks (u_track), the tracks are added to lost_tracks list and are marked lost
'''Deal with unconfirmed tracks, usually tracks with only one beginning frame'''
detections = [detections[i] for i in u_detection]
dists = matching.iou_distance(unconfirmed, detections)
matches, u_unconfirmed, u_detection = matching.linear_assignment(
dists, thresh=0.7)
for itracked, idet in matches:
unconfirmed[itracked].update(detections[idet], self.frame_id)
activated_starcks.append(unconfirmed[itracked])
# The tracks which are yet not matched
for it in u_unconfirmed:
track = unconfirmed[it]
track.mark_removed()
removed_stracks.append(track)
# after all these confirmation steps, if a new detection is found, it is initialized for a new track
""" Step 4: Init new stracks"""
for inew in u_detection:
track = detections[inew]
if track.score < self.det_thresh:
continue
track.activate(self.kalman_filter, self.frame_id)
activated_starcks.append(track)
""" Step 5: Update state"""
# If the tracks are lost for more frames than the threshold number, the tracks are removed.
for track in self.lost_stracks:
if self.frame_id - track.end_frame > self.max_time_lost:
track.mark_removed()
removed_stracks.append(track)
# print('Remained match {} s'.format(t4-t3))
# Update the self.tracked_stracks and self.lost_stracks using the updates in this step.
self.tracked_stracks = [
t for t in self.tracked_stracks if t.state == TrackState.Tracked
]
self.tracked_stracks = joint_stracks(self.tracked_stracks,
activated_starcks)
self.tracked_stracks = joint_stracks(self.tracked_stracks,
refind_stracks)
# self.lost_stracks = [t for t in self.lost_stracks if t.state == TrackState.Lost] # type: list[STrack]
self.lost_stracks = sub_stracks(self.lost_stracks,
self.tracked_stracks)
self.lost_stracks.extend(lost_stracks)
self.lost_stracks = sub_stracks(self.lost_stracks,
self.removed_stracks)
self.removed_stracks.extend(removed_stracks)
self.tracked_stracks, self.lost_stracks = remove_duplicate_stracks(
self.tracked_stracks, self.lost_stracks)
# get scores of lost tracks
output_stracks = [
track for track in self.tracked_stracks if track.is_activated
]
logger.debug('===========Frame {}=========='.format(self.frame_id))
logger.debug('Activated: {}'.format(
[track.track_id for track in activated_starcks]))
logger.debug('Refind: {}'.format(
[track.track_id for track in refind_stracks]))
logger.debug('Lost: {}'.format(
[track.track_id for track in lost_stracks]))
logger.debug('Removed: {}'.format(
[track.track_id for track in removed_stracks]))
# print('Final {} s'.format(t5-t4))
return output_stracks
