def __init__(self): self.t = None self.last_frame = None self.track_window = None self.s = ROSImageStream(sub_topic_name=args.i) self.p = ROSImageStream(pub_topic_name=args.o) self.pre_tracker = PreTracker() self.pre_tracker.set_color_range(((0,0,0),(10,10,10))) self.pre_tracker.set_pre_track_callbacks([self.pre_tracker.pre_track_color_range])
def start( self, loop=True ): if self.in_frames_topic is not None and self.out_frames_topic is not None: self.pub = ROSImageStream(pub_topic_name=self.out_frames_topic) self.sub = ROSImageStream(sub_topic_name=self.in_frames_topic) self.sub.img_stream_subscribe( callback=self.publish_frame, loop=False ) self.multi_tracker_receiver.start(loop=loop)
class TrackerExample: def __init__(self): self.t = None self.last_frame = None self.track_window = None self.s = ROSImageStream(sub_topic_name=args.i) self.p = ROSImageStream(pub_topic_name=args.o) self.pre_tracker = PreTracker() self.pre_tracker.set_color_range(((0,0,0),(10,10,10))) self.pre_tracker.set_pre_track_callbacks([self.pre_tracker.pre_track_color_range]) def start(self): def callback(frame): if self.t == None: seg = self.pre_tracker.pre_track(frame) cnts, _ = cv2.findContours( seg, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE ) max_cnt = max(cnts, key=lambda c: cv2.contourArea(c)) rectangle = cv2.boundingRect(max_cnt) self.t = CamshiftTracker(init_window=rectangle) r = cv2.rectangle(frame, (rectangle[0],rectangle[1]), (rectangle[2], rectangle[3]), (255, 0, 0),2) self.p.publish_single(r) self.last_frame = frame pre_tracked_frame = self.pre_tracker.pre_track(frame) ret, self.track_window = self.t.track(pre_tracked_frame) rectangle = cv2.rectangle(frame, (self.track_window[0], self.track_window[1]), (self.track_window[0] + self.track_window[2], self.track_window[1] + self.track_window[3]), (255,0,0), 2) self.p.publish_single(rectangle) self.s.img_stream_subscribe(callback)
def __init__( self, tracker_type=CamshiftTracker, # MeanshiftTracker, CamshiftTracker cnt_min_area=500, window_max_width=None, lost_ttl_attempts=10, max_tracked_objects=30, overlap_margin=25, allowed_directions=['n', 's', 'e', 'w'], pre_tracker=None, in_frames_topic="/img_stream", out_events_topic="/multi_tracker/global_events", out_status_topic="/multi_tracker/status", out_window_topic="/multi_tracker/window", in_kalman_topic="/multi_tracker/pos", out_frames_topic="/multi_tracker/tracked_frames"): self.tracker_type = tracker_type self.cnt_min_area = cnt_min_area self.window_max_width = window_max_width self.lost_ttl_attempts = lost_ttl_attempts self.max_tracked_objects = max_tracked_objects self.overlap_margin = overlap_margin self.allowed_directions = allowed_directions if pre_tracker == None: self.pre_tracker = PreTracker( ) #processed_frame_topic="/pre_tracked_frame") self.pre_tracker.set_color_range(color_range=((0, 0, 0), (20, 20, 20))) self.pre_tracker.set_median_blur(7) self.pre_tracker.set_pre_track_callbacks([ self.pre_tracker.pre_track_color_range, self.pre_tracker.pre_track_median_blur ]) else: self.pre_tracker = pre_tracker self.in_frames_topic = in_frames_topic self.out_events_topic = out_events_topic self.out_status_topic = out_status_topic self.out_window_topic = out_window_topic self.in_kalman_topic = in_kalman_topic self.out_frames_topic = out_frames_topic self.frames_img_stream = ROSImageStream( sub_topic_name=self.in_frames_topic) self.kalman_sub = rospy.Subscriber(self.in_kalman_topic, Window, callback=self.kalman_pos_callback, queue_size=self.max_tracked_objects) if self.out_frames_topic != None: self.out_frames_pub = ROSImageStream( pub_topic_name=self.out_frames_topic) else: selt.out_frames_pub = None self.tracked_objects = TrackedObjects( out_status_topic=self.out_status_topic, out_window_topic=self.out_window_topic, out_events_topic=self.out_events_topic, max_tracked_objects=self.max_tracked_objects, max_ttl=self.lost_ttl_attempts) self.trackers = {} # {"id": tracker_type}
class MultiTracker: def __init__( self, tracker_type=CamshiftTracker, # MeanshiftTracker, CamshiftTracker cnt_min_area=500, window_max_width=None, lost_ttl_attempts=10, max_tracked_objects=30, overlap_margin=25, allowed_directions=['n', 's', 'e', 'w'], pre_tracker=None, in_frames_topic="/img_stream", out_events_topic="/multi_tracker/global_events", out_status_topic="/multi_tracker/status", out_window_topic="/multi_tracker/window", in_kalman_topic="/multi_tracker/pos", out_frames_topic="/multi_tracker/tracked_frames"): self.tracker_type = tracker_type self.cnt_min_area = cnt_min_area self.window_max_width = window_max_width self.lost_ttl_attempts = lost_ttl_attempts self.max_tracked_objects = max_tracked_objects self.overlap_margin = overlap_margin self.allowed_directions = allowed_directions if pre_tracker == None: self.pre_tracker = PreTracker( ) #processed_frame_topic="/pre_tracked_frame") self.pre_tracker.set_color_range(color_range=((0, 0, 0), (20, 20, 20))) self.pre_tracker.set_median_blur(7) self.pre_tracker.set_pre_track_callbacks([ self.pre_tracker.pre_track_color_range, self.pre_tracker.pre_track_median_blur ]) else: self.pre_tracker = pre_tracker self.in_frames_topic = in_frames_topic self.out_events_topic = out_events_topic self.out_status_topic = out_status_topic self.out_window_topic = out_window_topic self.in_kalman_topic = in_kalman_topic self.out_frames_topic = out_frames_topic self.frames_img_stream = ROSImageStream( sub_topic_name=self.in_frames_topic) self.kalman_sub = rospy.Subscriber(self.in_kalman_topic, Window, callback=self.kalman_pos_callback, queue_size=self.max_tracked_objects) if self.out_frames_topic != None: self.out_frames_pub = ROSImageStream( pub_topic_name=self.out_frames_topic) else: selt.out_frames_pub = None self.tracked_objects = TrackedObjects( out_status_topic=self.out_status_topic, out_window_topic=self.out_window_topic, out_events_topic=self.out_events_topic, max_tracked_objects=self.max_tracked_objects, max_ttl=self.lost_ttl_attempts) self.trackers = {} # {"id": tracker_type} def start(self): self.frames_img_stream.img_stream_subscribe(self.apply_tracking) def apply_tracking(self, frame): pre_tracked_frame = self.pre_tracker.pre_track(frame) updated_tracked_frame = self.update_current_trackers(pre_tracked_frame) self.discover_new_objects(updated_tracked_frame) if self.out_frames_pub != None: self.publish_tracked_frame(frame) def update_current_trackers(self, pre_tracked_frame): for obj in self.tracked_objects: obj.acquire_lock() ret, track_window = self.trackers[obj.id].track(pre_tracked_frame) if self.window_max_width != None: if track_window[2] > self.window_max_width and track_window[ 3] > self.window_max_width: track_window = (track_window[0], track_window[1], self.window_max_width, self.window_max_width) elif track_window[2] > self.window_max_width: track_window = (track_window[0], track_window[1], self.window_max_width, track_window[3]) elif track_window[3] > self.window_max_width: track_window = (track_window[0], track_window[1], track_window[2], self.window_max_width) if obj.track_window != None and obj.last_window != None: last_vel = np.array([ obj.track_window[1] - obj.last_window[1], obj.track_window[0] - obj.last_window[0] ]) last_dir = None if abs(last_vel[0]) > abs(last_vel[1]): if last_vel[0] < 0: last_dir = 'n' else: last_dir = 's' else: if last_vel[1] < 0: last_dir = 'w' else: last_dir = 'e' current_vel = np.array([ track_window[1] - obj.track_window[1], track_window[0] - obj.track_window[0] ]) current_dir = None if abs(current_vel[0]) > abs(current_vel[1]): if current_vel[0] < 0: current_dir = 'n' else: current_dir = 's' else: if current_vel[1] < 0: current_dir = 'w' else: current_dir = 'e' if current_dir != last_dir: if (last_dir == "n" and current_dir == "s" or last_dir == "s" and current_dir == "n" or last_dir == "w" and current_dir == "e" or last_dir == "e" and current_dir == "w"): pre_tracked_frame = self.handle_lost_object( obj, pre_tracked_frame, immediate_remove=True) else: pre_tracked_frame = self.handle_lost_object( obj, pre_tracked_frame, immediate_remove=True) obj.release_lock() continue allowed_dir = False for direction in self.allowed_directions: if direction == current_dir: allowed_dir = True break if not allowed_dir: pre_tracked_frame = self.handle_lost_object( obj, pre_tracked_frame, immediate_remove=True) obj.release_lock() continue #if ( # obj.track_window[0] + obj.track_window[2] < 0 or # obj.track_window[0] > pre_tracked_frame.shape[1] or # obj.track_window[1] + obj.track_window[3] < 0 or # obj.track_window[1] > pre_tracked_frame.shape[0] #): # pre_tracked_frame = self.handle_lost_object( # obj, # pre_tracked_frame # ) # obj.release_lock() # continue x_lowerb, x_upperb, y_lowerb, y_upperb = self.get_slices( track_window, pre_tracked_frame) sub_frame = pre_tracked_frame[x_lowerb:x_upperb, y_lowerb:y_upperb] _, cnts, _ = cv2.findContours(sub_frame, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) if len(cnts) == 0: pre_tracked_frame = self.handle_lost_object( obj, pre_tracked_frame) obj.release_lock() continue max_cnt = cnts[0] max_cnt_area = cv2.contourArea(max_cnt) for cnt in cnts: cnt_area = cv2.contourArea(cnt) if cnt_area > max_cnt_area: max_cnt_area = cnt_area if max_cnt_area < self.cnt_min_area: pre_tracked_frame = self.handle_lost_object( obj, pre_tracked_frame) obj.release_lock() continue obj.supply_window(track_window) pre_tracked_frame = self.handle_ok_object(obj, pre_tracked_frame) obj.release_lock() return pre_tracked_frame def handle_lost_object(self, obj, pre_tracked_frame, immediate_remove=False): if immediate_remove: obj.ttl = 0 obj.set_lost() x_lowerb, x_upperb, y_lowerb, y_upperb = self.get_slices( obj.track_window, pre_tracked_frame) pre_tracked_frame[x_lowerb:x_upperb, y_lowerb:y_upperb] = np.full( pre_tracked_frame[x_lowerb:x_upperb, y_lowerb:y_upperb].shape, 0) return pre_tracked_frame def handle_ok_object(self, obj, pre_tracked_frame): obj.set_ok() x_lowerb, x_upperb, y_lowerb, y_upperb = self.get_slices( obj.track_window, pre_tracked_frame) pre_tracked_frame[x_lowerb:x_upperb, y_lowerb:y_upperb] = np.full( pre_tracked_frame[x_lowerb:x_upperb, y_lowerb:y_upperb].shape, 0) return pre_tracked_frame def get_slices(self, track_window, pre_tracked_frame): x_lowerb = track_window[1] - self.overlap_margin if x_lowerb < 0: x_lowerb = 0 x_upperb = track_window[1] + track_window[3] + self.overlap_margin if x_upperb > pre_tracked_frame.shape[0]: x_upperb = pre_tracked_frame.shape[0] y_lowerb = track_window[0] - self.overlap_margin if y_lowerb < 0: y_lowerb = 0 y_upperb = track_window[0] + track_window[2] + self.overlap_margin if y_upperb > pre_tracked_frame.shape[1]: y_upperb = pre_tracked_frame.shape[1] return x_lowerb, x_upperb, y_lowerb, y_upperb def discover_new_objects(self, updated_frame): _, cnts, _ = cv2.findContours(updated_frame, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) bounding_rects = [] for cnt in cnts: rect = cv2.boundingRect(cnt) if self.window_max_width != None: if rect[2] > self.window_max_width and rect[ 3] > self.window_max_width: rect = (rect[0], rect[1], self.window_max_width, self.window_max_width) elif rect[2] > self.window_max_width: rect = (rect[0], rect[1], self.window_max_width, rect[3]) elif rect[3] > self.window_max_width: rect = (rect[0], rect[1], rect[2], self.window_max_width) overlapping = False for b_rect in bounding_rects: if not (rect[1] + rect[2] < b_rect[1] or b_rect[1] + b_rect[2] < rect[1] or rect[0] + rect[3] < b_rect[0] or b_rect[0] + b_rect[3] < rect[0]): overlapping = True break if overlapping: continue cnt_area = cv2.contourArea(cnt) if cnt_area < self.cnt_min_area: continue bounding_rects.append(rect) key = self.tracked_objects.create(window=rect) if key != None: self.trackers[key] = self.tracker_type(init_window=rect) def kalman_pos_callback(self, msg): obj = self.tracked_objects.get_item(msg.id) if obj != None: obj.acquire_lock() if obj.status == "lost": obj.track_window = msg.window obj.release_lock() def publish_tracked_frame(self, frame): #frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR) for obj in self.tracked_objects.tracked_objects: obj.acquire_lock() if obj.status == "ok": frame = cv2.rectangle( frame, (obj.track_window[0], obj.track_window[1]), (obj.track_window[0] + obj.track_window[2], obj.track_window[1] + obj.track_window[3]), (255, 0, 0), thickness=3) elif obj.status == "lost": frame = cv2.rectangle( frame, (obj.track_window[0], obj.track_window[1]), (obj.track_window[0] + obj.track_window[2], obj.track_window[1] + obj.track_window[3]), (0, 0, 255), thickness=3) obj.release_lock() self.out_frames_pub.publish_single(frame)
class KalmanMultiTrackerReceiver(): class KalmanTrackerReceiver(RTKalmanFilter): def __init__( self, dt_prediction=0.05, prediction_topic=None, update_topic=None, state_est_topic="/rt_kalman_filter/state_est", start_pos=None ): self.dt_prediction = dt_prediction self.prediction_topic = prediction_topic self.update_topic = update_topic self.state_est_topic = state_est_topic self.sensor_bag = None self.sensor_topic = None self.pub_prediction = None self.pub_update = None self.pub_state_est = None x = None if start_pos != None: x = np.array([ [0.], #px [0.], #py [0.], #vx [0.]] #vy ) else: x = np.array([ [start_pos[0, 0]], [start_pos[1, 0]], [0.], [0.]] ) F = np.array([ [1., 0, self.dt_prediction, 0], [0., 1., 0., self.dt_prediction], [0., 0., 1., 0.], [0., 0., 0., 1.]] ) H = np.array([ [0., 1., 0., 0.], [1., 0., 0., 0.]] ) P = np.array([ [100000. ,0., 0., 0.], [0., 100000., 0., 0.], [0., 0., 100000., 0.], [0., 0., 0., 100000.]] ) R = np.array([ [1.5, 0], [0., 1.5]] ) Q = Q_discrete_white_noise(dim=4, var=0.13) self.init_filter( x, F, H, P, R, Q ) self.filter_lock = threading.Lock() self.update_thread = None self.prediction_thread = None def __init__( self, max_tracked_objects=30, dt_prediction=0.05, prediction_topic=None, update_topic=None, state_est_topic="/kalman_multi_tracker/state_est", in_events_topic="/multi_tracker/global_events", in_status_topic="/multi_tracker/status", in_window_topic="/multi_tracker/window", out_kalman_topic="/multi_tracker/pos", in_frames_topic=None, out_frames_topic=None ): self.max_tracked_objects = max_tracked_objects self.dt_prediction = dt_prediction self.prediction_topic = prediction_topic self.update_topic = update_topic self.state_est_topic = state_est_topic self.in_events_topic = in_events_topic self.in_status_topic = in_status_topic self.in_window_topic = in_window_topic self.out_kalman_topic = out_kalman_topic self.in_frames_topic = in_frames_topic self.out_frames_topic = out_frames_topic self.kalman_filters = {} # {"key":KalmanTrackerReceiver} self.kalman_filters_lock = threading.Lock() self.multi_tracker_receiver = MultiTrackerReceiver( max_tracked_objects=self.max_tracked_objects, event_callback=self.event_callback, status_callback=self.status_callback, position_callback=self.position_callback, in_events_topic=self.in_events_topic, in_status_topic=self.in_status_topic, in_window_topic=self.in_window_topic, out_kalman_topic=self.out_kalman_topic ) def start( self, loop=True ): if self.in_frames_topic is not None and self.out_frames_topic is not None: self.pub = ROSImageStream(pub_topic_name=self.out_frames_topic) self.sub = ROSImageStream(sub_topic_name=self.in_frames_topic) self.sub.img_stream_subscribe( callback=self.publish_frame, loop=False ) self.multi_tracker_receiver.start(loop=loop) def publish_frame( self, frame ): self.kalman_filters_lock.acquire() for k in self.kalman_filters.keys(): frame = cv2.circle( frame, (int(self.kalman_filters[k].filter.x[0,0]), int(self.kalman_filters[k].filter.x[1,0])), 10, (255,0,0) ) self.pub.publish_single(frame) self.kalman_filters_lock.release() def event_callback( self, key, event ): self.kalman_filters_lock.acquire() if event == "+": # self.kalman_filters[key] = None #self.kalman_filters[key] = self.KalmanTrackerReceiver( # dt_prediction=self.dt_prediction, # prediction_topic=self.prediction_topic, # update_topic=self.update_topic, # state_est_topic=self.state_est_topic #) #self.kalman_filters[key].start() #if key in self.kalman_filters.keys(): # self.kalman_filters[key].stop() # self.kalman_filters[key] = None pass elif event == "-": if key in self.kalman_filters.keys(): self.kalman_filters[key].stop() del self.kalman_filters[key] self.kalman_filters_lock.release() def status_callback( self, key, status ): self.kalman_filters_lock.acquire() if status == "lost": if key in self.kalman_filters.keys(): callbacks = { "prediction": { "callback": self.multi_tracker_receiver.supply_kalman_estimate, "arg": key } } self.kalman_filters[key].set_callbacks(callbacks) elif status == "ok" or status == "removed": if key in self.kalman_filters.keys(): self.kalman_filters[key].set_callbacks(None) self.kalman_filters_lock.release() def position_callback( self, key, position ): self.kalman_filters_lock.acquire() if key not in self.kalman_filters.keys(): self.kalman_filters[key] = self.KalmanTrackerReceiver( dt_prediction=self.dt_prediction, prediction_topic=self.prediction_topic, update_topic=self.update_topic, state_est_topic=self.state_est_topic, start_pos=position ) self.kalman_filters[key].start() else: self.kalman_filters[key].update_step(position) self.kalman_filters_lock.release()