def heuristictype_detection(self,track_trajectory,mov_id,orig_pred): if orig_pred !='car': return orig_pred else: if str(mov_id) in['1','6']: return self.person_or_motorcycle[0] if tlbrs_to_mean_area(track_trajectory)<=10000 else orig_pred elif str(mov_id) in ['3','4']: return self.person_or_motorcycle[0] if tlbrs_to_mean_area(track_trajectory)<=8500 else orig_pred elif str(mov_id) in ['2','5']: return self.person_or_motorcycle[0] if tlbrs_to_mean_area(track_trajectory)<=3500 else orig_pred
def update(self, im_blob, img0): self.frame_id += 1 activated_starcks = [] refind_stracks = [] lost_stracks = [] removed_stracks = [] width = img0.shape[1] height = img0.shape[0] init_polygon=self.polygon2 if self.two_polygon_system and self.frame_id>= self.warmup_frame else self.polygon two_wheel_polygon=self.polygon four_wheel_polygon=self.polygon virtual_polygon=self.virtual_polygon huge_box_thres=150 bbox=[] score=[] types=[] huge_vehicles=[] ''' Step 1: Network forward, get detections & embeddings''' if self.opt.detection_model=='Efficient': with torch.no_grad(): ori_imgs, framed_imgs, framed_metas = preprocess([img0], max_size=self.input_size) device = torch.device('cuda:0') x = torch.stack([torch.from_numpy(fi).to(device) for fi in framed_imgs], 0) x = x.to(torch.float32 ).permute(0, 3, 1, 2) features, regression, classification, anchors = self.detetection_model(x) regressBoxes = BBoxTransform() clipBoxes = ClipBoxes() out = postprocess(x, anchors, regression, classification, regressBoxes, clipBoxes, self.opt.det_thres, self.opt.nms_thres) out = invert_affine(framed_metas, out) bbox=[] score=[] types=[] huge_vehicles=[] for j in range(len(out[0]['rois'])): obj = self.obj_list[out[0]['class_ids'][j]] if obj in self.obj_interest: x1, y1, x2, y2 = out[0]['rois'][j].astype(np.int) #bike,bicycle if (y1+y2)/2>height/2 and float(out[0]['scores'][j])<=0.25: continue if (y2-y1)*(x2-x1)<=200: continue if obj not in self.person_or_motorcycle and float(out[0]['scores'][j])>=0.3: bbox.append([x1, y1, x2, y2]) score.append( float(out[0]['scores'][j])) types.append(obj) huge_vehicles.append(False if ((y2-y1)<=huge_box_thres and (x2-x1)<=180) else True ) elif obj in self.person_or_motorcycle: #['bicycle', 'motorcycle'] bbox.append([x1, y1, x2, y2]) score.append( float(out[0]['scores'][j])) types.append(obj) huge_vehicles.append(False) elif self.opt.detection_model=='FasterRcnn': predictions= self.detetection_model.compute_prediction(img0) top_predictions=self.detetection_model.select_top_predictions(predictions) scores = top_predictions.get_field("scores").tolist() labels = top_predictions.get_field("labels").tolist() labels = [self.detetection_model.CATEGORIES[i] for i in labels] boxes = top_predictions.bbox.tolist() for j in range(len(labels)): obj = labels[j] if obj in self.obj_interest: x1, y1, x2, y2 = boxes[j] if (y1+y2)/2>0.5*height and float(scores[j])<=0.25: continue if (y2-y1)*(x2-x1)<=200: continue if obj not in self.person_or_motorcycle and float(scores[j])>=0.3: bbox.append([x1, y1, x2, y2]) score.append( float(scores[j])) types.append(obj) huge_vehicles.append(False if ((y2-y1)<=huge_box_thres and (x2-x1)<=180) else True ) elif obj in self.person_or_motorcycle and float(scores[j])>=self.opt.det_thres: #['bicycle', 'motorcycle'] bbox.append([x1, y1, x2, y2]) score.append( float(scores[j])) types.append(obj) huge_vehicles.append(False) # vis # print(len(bbox)) # print(img0.shape) # print(self.polygon) # for i in range(len(bbox)): # bb = bbox[i] # cv2.rectangle(img0, (bb[0], bb[1]), # (bb[2], bb[3]), # (0, 255, 0), 2) # cv2.polylines(img0,[np.asarray(self.polygon)],True,(0,255,255)) # cv2.imshow('dets', img0) # cv2.waitKey(0) if len(bbox) > 0: '''Detections''' detections = [STrack(STrack.tlbr_to_tlwh(tlbr), sco, clas, 30,huge_vehicle=hv) for (tlbr, sco,clas,hv) in zip(bbox,score,types,huge_vehicles)] else: detections = [] detections_plot=copy.deepcopy(detections) ''' Add newly detected tracklets to tracked_stracks''' unconfirmed = [] tracked_stracks = [] # type: list[STrack] for track in self.tracked_stracks: if not track.is_activated: unconfirmed.append(track) else: tracked_stracks.append(track) ''' Step 2: First association, with gating distance''' strack_pool,lost_map_tracks = joint_stracks(tracked_stracks, self.lost_stracks) # Predict the current location with KF #for strack in strack_pool: #strack.predict() STrack.multi_predict(strack_pool) #dists = matching.embedding_distance(strack_pool, detections) detections=heuristic_occlusion_detection(detections) match_thres=100 dists=np.zeros(shape=(len(strack_pool),len(detections))) dists = matching.gate_cost_matrix3(self.kalman_filter, dists, strack_pool, detections,type_diff=True) #dists = matching.fuse_motion(self.opt,self.kalman_filter, dists, strack_pool, detections,lost_map=lost_map_tracks,occlusion_map=occlusion_map,thres=match_thres) matches, u_track, u_detection = matching.linear_assignment(dists, thresh=match_thres) for itracked, idet in matches: track = strack_pool[itracked] det = detections[idet] if track.state == TrackState.Tracked: track.update(detections[idet], self.frame_id) activated_starcks.append(track) else: track.re_activate(det, self.frame_id, new_id=False) refind_stracks.append(track) ''' Step 3: Second association, with IOU''' detections = [detections[i] for i in u_detection] r_tracked_stracks = [strack_pool[i] for i in u_track if strack_pool[i].state == TrackState.Tracked] dists = matching.iou_distance(r_tracked_stracks, detections) matches, u_track, u_detection = matching.linear_assignment(dists, thresh=0.5) 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) ''' ''' for it in u_track: track = r_tracked_stracks[it] if not track.state == TrackState.Lost: track.mark_lost() lost_stracks.append(track) '''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.6) for itracked, idet in matches: unconfirmed[itracked].update(detections[idet], self.frame_id) activated_starcks.append(unconfirmed[itracked]) for it in u_unconfirmed: track = unconfirmed[it] track.mark_removed() removed_stracks.append(track) """ Step 4: Init new stracks""" for inew in u_detection: track = detections[inew] track_init_polygon=init_polygon if not track.huge_vehicle else virtual_polygon if track.score < self.det_thresh or track.occlusion_status==True or check_bbox_outside_polygon(track_init_polygon,track.tlbr): continue # track_types=self.person_or_motorcycle[0] if tlbrs_to_mean_area(track.track_trajectory) <=1500 else track.infer_type() if self.frame_id>=1 and not check_bbox_inside_polygon(track_init_polygon,track.tlbr):#and track_types in self.person_or_motorcycle #person, motorcycle continue track.activate(self.kalman_filter, self.frame_id) activated_starcks.append(track) """ Step 5: Update state and getting out of interest tracklet if have""" out_of_polygon_tracklet=[] refind_stracks_copy=[] activated_starcks_copy=[] for idx,current_tracked_tracks in enumerate([refind_stracks,activated_starcks]) :# for track in current_tracked_tracks: if tlbrs_to_mean_area(track.track_trajectory) <=1000 : track_type= self.person_or_motorcycle[0] #person else: track_type=track.infer_type() if track_type in self.person_or_motorcycle: out_polygon=two_wheel_polygon p_type='two_wheel' else: out_polygon=four_wheel_polygon #if not track.huge_vehicle else virtual_polygon p_type='four_wheel' if check_bbox_outside_polygon(out_polygon,track.tlbr) : track.mark_removed() removed_stracks.append(track) if ((len(track.track_frames)>=4 and self.frame_id <=5) or (len(track.track_frames)>=5 and self.frame_id>=self.warmup_frame+5)) and idx==1:########## 4 is confident number of frame track_center=[ [(x[0]+x[2])/2,(x[1]+x[3])/2] for x in track.track_trajectory] # movement_id=counting_moi(self.paths,[(track_center[0],track_center[-1])])[0] # movement_id=self.heusristic_mov_refinement(track,movement_id) movement_id=self.heusristic_mov_detection(track,track_center) frame_id=self.frame_id out_of_polygon_tracklet.append((frame_id,track.track_id,track_type,movement_id)) else: refind_stracks_copy.append(track) if idx ==0 else activated_starcks_copy.append(track) refind_stracks=refind_stracks_copy activated_starcks=activated_starcks_copy lost_stracks_copy=[] for track in lost_stracks: if tlbrs_to_mean_area(track.track_trajectory) <=1000 : track_type= self.person_or_motorcycle[0] #person else: track_type=track.infer_type() if track_type in self.person_or_motorcycle: out_polygon=two_wheel_polygon p_type='two_wheel' else: out_polygon=four_wheel_polygon p_type='four_wheel' if check_bbox_intersect_or_outside_polygon(out_polygon,track.tlbr) : track.mark_removed() removed_stracks.append(track) if ((len(track.track_frames)>=4 and self.frame_id <=5) or (len(track.track_frames)>=6 and self.frame_id>=self.warmup_frame+5)): track_center=[ [(x[0]+x[2])/2,(x[1]+x[3])/2] for x in track.track_trajectory] # movement_id=counting_moi(self.paths,[(track_center[0],track_center[-1])])[0] # movement_id=self.heusristic_mov_refinement(track,movement_id) movement_id=self.heusristic_mov_detection(track,track_center) # line_interest=self.line1 if str(movement_id)=='1' else self.line2 # out_direction='bottom' if str(movement_id)=='1' else 'up' frame_id=self.frame_id if track_type in self.person_or_motorcycle else self.frame_id+5#kalman_predict_out_line(track,line_interest,out_direction) out_of_polygon_tracklet.append((frame_id,track.track_id,track_type,movement_id)) else: lost_stracks_copy.append(track) lost_stracks=lost_stracks_copy for track in self.lost_stracks: if self.frame_id - track.end_frame > self.max_time_lost : track.mark_removed() removed_stracks.append(track) #Remove out of screen tracklet elif track.tlwh[0]+track.tlwh[2]//2>width or track.tlwh[1]+track.tlwh[3]//2>height or min(track.tlwh[0]+track.tlwh[2]//2,track.tlwh[1]+track.tlwh[3]//2)<0: track.num_out_frame+=1 if track.num_out_frame>STrack.out_of_frame_patience: track.mark_removed() removed_stracks.append(track) # print('Ramained match {} s'.format(t4-t3)) # print(out_of_polygon_tracklet) 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 = 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) #self.merge_track() 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])) return output_stracks,detections_plot,out_of_polygon_tracklet