def initTrack(self): if (self.track): self.progressBar.setValue(20) if self.options["track"]: if self.options["tracker"] == "deep_sort": from deep_sort import generate_detections from deep_sort.deep_sort import nn_matching from deep_sort.deep_sort.tracker import Tracker self.progressBar.setValue(50) metric = nn_matching.NearestNeighborDistanceMetric( "cosine", 0.2, 100) self.tracker = Tracker(metric) self.encoder = generate_detections.create_box_encoder( os.path.abspath( "deep_sort/resources/networks/mars-small128.ckpt-68577" )) elif self.options["tracker"] == "sort": from sort.sort import Sort self.encoder = None self.tracker = Sort() self.progressBar.setValue(50) if self.options["BK_MOG"] and self.options["track"]: fgbg = cv2.bgsegm.createBackgroundSubtractorMOG() self.progressBar.setValue(60) self.initTFNet() else: self.initTFNet()
def __init__(self): #loading this encoder is slow, should be done only once. #self.encoder = generate_detections.create_box_encoder("deep_sort/resources/networks/mars-small128.ckpt-68577") self.encoder = vehicle_encoder('weights/vehicle_vgg.h5') self.metric = nn_matching.NearestNeighborDistanceMetric( "cosine", .5, 100) self.tracker = Tracker(self.metric)
def __init__(self, debug=0): self.debug = debug self.frameCount = 0 self.timePrev = time.time() #region DEEPSORT metric = nn_matching.NearestNeighborDistanceMetric( 'cosine', matching_threshold=0.2, budget=100) self.tracker = Tracker(metric) absFilePath = os.path.abspath(__file__) fileDir = os.path.dirname(absFilePath) filePathModel = os.path.join( fileDir, 'deep_sort/resources/networks/mars-small128.pb') self.encoder = gdet.create_box_encoder(filePathModel, batch_size=1) #endregion self.detectionAndId = [] self.subscriberImageDetections = rospy.Subscriber( 'yolo_detector/output/compresseddetections', CompressedImageAndBoundingBoxes, self.callback, queue_size=1) self.publisherDetectionID = rospy.Publisher('yact/output/detectionids', DetectionAndID, queue_size=1)
def __init__(self, classes, tracker='sort'): self.ttype = tracker self.classes = classes if tracker == 'deep_sort': from deep_sort import generate_detections from deep_sort.deep_sort import nn_matching from deep_sort.deep_sort.tracker import Tracker metric = nn_matching.NearestNeighborDistanceMetric( "cosine", 0.2, 100) #param self.nms_max_overlap = 0.1 #param model_path = os.path.join(WORK_DIR, MODEL_DIR, "mars-small128.ckpt-68577") self.encoder = generate_detections.create_box_encoder(model_path) self.tracker = Tracker(metric) from deep_sort.application_util import preprocessing as prep from deep_sort.deep_sort.detection import Detection self.prep = prep self.Detection = Detection elif tracker == 'sort': from sort.sort import Sort self.tracker = Sort() self.trackers = {}
def __init__(self, wt_path=None): #loading this encoder is slow, should be done only once. #self.encoder = generate_detections.create_box_encoder("deep_sort/resources/networks/mars-small128.ckpt-68577") if wt_path is not None: self.encoder = torch.load(wt_path) else: dirname = os.path.dirname(os.path.abspath(__file__)) filename = os.path.join(dirname, 'ckpts/model640.pt') self.encoder = torch.load(filename) self.encoder = self.encoder.cuda() self.encoder = self.encoder.eval() print("Deep sort model loaded") self.metric = nn_matching.NearestNeighborDistanceMetric( "cosine", .5, 100) self.tracker = Tracker(self.metric) self.gaussian_mask = get_gaussian_mask().cuda() self.transforms = torchvision.transforms.Compose([ \ torchvision.transforms.ToPILImage(),\ torchvision.transforms.Resize((128,128)),\ torchvision.transforms.ToTensor()])
def gui(self): source = self.FLAGS.demo SaveVideo = self.FLAGS.saveVideo if self.FLAGS.track: if self.FLAGS.tracker == "deep_sort": from deep_sort import generate_detections from deep_sort.deep_sort import nn_matching from deep_sort.deep_sort.tracker import Tracker metric = nn_matching.NearestNeighborDistanceMetric( "cosine", 0.2, 100) tracker = Tracker(metric) encoder = generate_detections.create_box_encoder( os.path.abspath( "deep_sort/resources/networks/mars-small128.ckpt-68577")) elif self.FLAGS.tracker == "sort": from sort.sort import Sort encoder = None tracker = Sort() if self.FLAGS.BK_MOG and self.FLAGS.track: fgbg = cv2.bgsegm.createBackgroundSubtractorMOG() if self.FLAGS.csv: f = open('{}.csv'.format(file), 'w') writer = csv.writer(f, delimiter=',') writer.writerow(['frame_id', 'track_id', 'x', 'y', 'w', 'h']) f.flush() else: f = None writer = None App(tkinter.Tk(), "Tkinter and OpenCV", 0, tracker, encoder)
def __init__(self): self.metric = nn_matching.NearestNeighborDistanceMetric( "cosine", 0.2, 100) self.tracker = Tracker(self.metric) self.encoder = generate_detections.create_box_encoder( os.path.abspath( "deep_sort/resources/networks/mars-small128.ckpt-68577"))
def __init__(self, encoderPath, applyMask=False): self.metric = nn_matching.NearestNeighborDistanceMetric( "cosine", 0.9, 100) self.tracker = Tracker(self.metric, max_iou_distance=0.9, max_age=50, n_init=3, _lambda=0.3) self.encoder = generate_detections.create_box_encoder( encoderPath, applyMask=applyMask) self.applyMask = applyMask
def __init__(self, deep_sort_feature_model, deep_sort_max_cosine_distance, deep_sort_nn_budget, per_process_gpu_mem_fraction, app_gpu): metric = nn_matching.NearestNeighborDistanceMetric( "cosine", deep_sort_max_cosine_distance, deep_sort_nn_budget) self.tracker = Tracker(metric) self.encoder = DeepSortTracker.create_box_encoder( deep_sort_feature_model, batch_size=32, per_process_gpu_mem_fraction=0.1, app_gpu=0)
def __init__(self, args): self.tfnet = None self.image_publisher = rospy.Publisher(args.output, Image, queue_size=1) self.subscriber = rospy.Subscriber(args.input, Image, self.callback, queue_size = 1, buff_size=2**24) self.subscriber = rospy.Subscriber("/cctv_info", String, self.exit) metric = nn_matching.NearestNeighborDistanceMetric( "cosine", 0.2, 100) self.tracker = Tracker(metric) self.encoder = generate_detections.create_box_encoder( os.path.abspath("deep_sort/resources/networks/mars-small128.ckpt-68577")) options = {"model": "darkflow/cfg/yolo.cfg", "load": "darkflow/bin/yolo.weights", "threshold": 0.1, "track": True, "trackObj": ["person"], "BK_MOG": True, "tracker": "deep_sort", "csv": False} self.tfnet = TFNet(options)
def __init__(self, model_folder, max_age, max_distance=0.1, nn_budget=None, nms_max_overlap=1.0, n_init=3): # Definition of the parameters self.max_distance = max_distance self.nn_budget = nn_budget self.nms_max_overlap = nms_max_overlap # deep_sort self.encoder = gdet.create_box_encoder(model_folder, batch_size=1) #self.metric = nn_matching.NearestNeighborDistanceMetric("cosine", 0.5 , nn_budget) #max_cosine_distance self.metric = nn_matching.NearestNeighborDistanceMetric( "euclidean", max_distance, nn_budget) self.tracker = Tracker(self.metric, max_age=max_age, n_init=n_init)
def predict(self, boxes, confidences, classids, crop_features): show_box = [] show_confidence = [] show_id = [] show_name = [] indices = [] for index in range(len(self.labels)): if type(self.trackers[index]) is list: if len(classids[index]) != 0: self.trackers[index] = Tracker( nn_matching.NearestNeighborDistanceMetric( "cosine", self.max_cosine_distance, self.nn_budget)) else: continue detections = [ dt.Detection(bbox, score, features, class_id) for bbox, score, features, class_id in zip( np.array(boxes[index]), np.array(confidences[index]), crop_features[index], classids[index]) ] self.trackers[index].predict() self.trackers[index].update(detections) for track in self.trackers[index].tracks: if not track.is_confirmed() or track.time_since_update > 1: continue show_box.append(track.to_tlwh()) show_confidence.append(track.confidence) show_id.append(track.track_id) show_name.append(track.get_name()) indices = preprocessing.non_max_suppression( np.array(show_box), self.nms_threshold, np.array(show_confidence)) return [show_box[i] for i in indices], [show_id[i] for i in indices ], [show_name[i] for i in indices]
def __init__(self, wt_path='ckpts/cos_metric_net.pt', wt_type='cosine'): if wt_type == 'cosine': exclude_keys = ['weights', 'scale'] # these params useless in eval stage self.encoder = CosineMetricNet(num_classes=-1, add_logits=False).cuda() try: ckpt: collections.OrderedDict = torch.load(wt_path) # type: ckpt['model_state_dict']: dict ckpt['model_state_dict'] = { k: v for k, v in ckpt['model_state_dict'].items() if k not in exclude_keys } self.encoder.load_state_dict(ckpt['model_state_dict'], strict=False) except KeyError as e: s = "Model loaded(%s) is not compatible with the definition, please check!" % wt_path raise KeyError(s) from e elif wt_type == 'siamese': self.encoder = torch.load(wt_path) else: raise NotImplementedError self.encoder = self.encoder.eval() print("Deep sort model loaded") self.metric = nn_matching.NearestNeighborDistanceMetric( "cosine", .5, 100) self.tracker = Tracker(self.metric) self.transforms = torchvision.transforms.Compose([ torchvision.transforms.ToPILImage(), torchvision.transforms.Resize((128, 64)), torchvision.transforms.ToTensor(), torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ])
def __init__(self,wt_path=None, max_age = 180, matching_threshold = 0.5, max_iou_distance = 0.7): #loading this encoder is slow, should be done only once. #self.encoder = generate_detections.create_box_encoder("deep_sort/resources/networks/mars-small128.ckpt-68577") if wt_path is not None: self.encoder = torch.load(wt_path) else: self.encoder = torch.load('ckpts/model640.pt') self.encoder = self.encoder.cuda() self.encoder = self.encoder.eval() print("Deep sort model loaded") self.metric = nn_matching.NearestNeighborDistanceMetric("cosine",matching_threshold = matching_threshold , budget = 100) self.tracker= Tracker(self.metric,max_iou_distance = max_iou_distance, max_age = max_age) self.gaussian_mask = get_gaussian_mask().cuda() self.transforms = torchvision.transforms.Compose([ \ torchvision.transforms.ToPILImage(),\ torchvision.transforms.Resize((128,128)),\ torchvision.transforms.ToTensor()])
def __init__( self, wt_path="/gdrive/MyDrive/car_tracking/car_tracking/object_tracking/ckpts/model640.pt" ): #loading this encoder is slow, should be done only once. #self.encoder = generate_detections.create_box_encoder("deep_sort/resources/networks/mars-small128.ckpt-68577") self.encoder = torch.load(wt_path) self.encoder = self.encoder.cuda() self.encoder = self.encoder.eval() print("Deep sort model loaded from path: ", wt_path) self.metric = nn_matching.NearestNeighborDistanceMetric( "cosine", .5, 100) self.tracker = Tracker(self.metric) self.gaussian_mask = get_gaussian_mask().cuda() self.transforms = torchvision.transforms.Compose([ \ torchvision.transforms.ToPILImage(),\ torchvision.transforms.Resize((128,128)),\ torchvision.transforms.ToTensor()])
def camera(self): file = self.FLAGS.demo SaveVideo = self.FLAGS.saveVideo if self.FLAGS.track: if self.FLAGS.tracker == "deep_sort": from deep_sort import generate_detections from deep_sort.deep_sort import nn_matching from deep_sort.deep_sort.tracker import Tracker metric = nn_matching.NearestNeighborDistanceMetric( "cosine", 0.2, 100) tracker = Tracker(metric) encoder = generate_detections.create_box_encoder( os.path.abspath( "deep_sort/resources/networks/mars-small128.ckpt-68577")) elif self.FLAGS.tracker == "sort": from sort.sort import Sort encoder = None tracker = Sort() if self.FLAGS.BK_MOG and self.FLAGS.track: fgbg = cv2.createBackgroundSubtractorMOG2() if file == 'camera': file = 0 else: assert os.path.isfile(file), \ 'file {} does not exist'.format(file) vid = imageio.get_reader(file, 'ffmpeg') #cv2.VideoCapture(file) if file == 0: self.say('Press [ESC] to quit video') #assert camera.isOpened(), \ #'Cannot capture source' if self.FLAGS.csv: f = open('{}.csv'.format(file), 'w') writer = csv.writer(f, delimiter=',') writer.writerow(['frame_id', 'track_id', 'x', 'y', 'w', 'h']) f.flush() else: f = None writer = None # buffers for demo in batch buffer_inp = list() buffer_pre = list() elapsed = 0 start = timer() self.say('Press [ESC] to quit demo') #postprocessed = [] # Loop through frames n = 0 plt.ion() fig = plt.figure() ax = plt.gca() frame = vid.get_data(0) img_artist = ax.imshow(frame) for num in range(1, 20000): try: frame = vid.get_data(num) print(num) except: break elapsed += 1 #_, frame = camera.read() if frame is None: print('\nEnd of Video') break if self.FLAGS.skip != n: n += 1 continue n = 0 if self.FLAGS.BK_MOG and self.FLAGS.track: fgmask = fgbg.apply(frame) else: fgmask = None preprocessed = self.framework.preprocess(frame) buffer_inp.append(frame) buffer_pre.append(preprocessed) # Only process and imshow when queue is full if elapsed % self.FLAGS.queue == 0: feed_dict = {self.inp: buffer_pre} net_out = self.sess.run(self.out, feed_dict) for img, single_out in zip(buffer_inp, net_out): if not self.FLAGS.track: postprocessed = self.framework.postprocess(single_out, img) else: postprocessed = self.framework.postprocess( single_out, img, frame_id=elapsed, csv_file=f, csv=writer, mask=fgmask, encoder=encoder, tracker=tracker) if self.FLAGS.display: #cv2.imshow('', postprocessed) img_artist.set_data(postprocessed) plt.show() plt.pause(0.00001) # Clear Buffers buffer_inp = list() buffer_pre = list() if elapsed % 5 == 0: sys.stdout.write('\r') sys.stdout.write('{0:3.3f} FPS'.format(elapsed / (timer() - start))) sys.stdout.flush() sys.stdout.write('\n') if self.FLAGS.csv: f.close()
def __init__(self): self.metric = nn_matching.NearestNeighborDistanceMetric( "cka", .5, 30) #metric, matching_threshold, budget self.tracker = Tracker(self.metric)
personDetector = PersonDetector(args) deepSort = False sort = True # SORT mot_tracker = Sort() # DEEP SORT PARAMS # Definition of the parameters max_cosine_distance = 0.3 nn_budget = None model_filename = 'deep_sort/models/mars-small128.pb' encoder = gdet.create_box_encoder(model_filename, batch_size=1) metric = nn_matching.NearestNeighborDistanceMetric( "cosine", max_cosine_distance, nn_budget) tracker = Tracker(metric) # Open Webcam video_capturer = cv2.VideoCapture(0) video_capturer.set(cv2.CAP_PROP_FRAME_WIDTH, 720) video_capturer.set(cv2.CAP_PROP_FRAME_HEIGHT, 480) infer_time = [] videoSize = (video_capturer.get(3), video_capturer.get(4)) while video_capturer.isOpened(): _, frame = video_capturer.read() bboxes = personDetector.Detect(frame) # Draw detections for bbox in bboxes:
def vis_detections_video(im, class_name, dets, csv_file, csv, frame_id, thresh=0.5): """Draw detected bounding boxes.""" nms_max_overlap = 0.6 metric = nn_matching.NearestNeighborDistanceMetric("cosine", 0.2, 100) tracker = Tracker(metric) detections = [] scores = [] h, w, _ = im.shape thick = int((h + w) // 300) inds = np.where(dets[:, -1] >= thresh)[0] if len(inds) == 0: return im for i in inds: scores.append(dets[i, -1]) for i in inds: bbox = dets[i, :4] boxResults = process_box(bbox, scores, h, w, thresh) if boxResults is None: continue left, right, top, bot, mess, max_indx, confidence = boxResults detections.append( np.array([left, top, right - left, bot - top]).astype(np.float64)) scores.append(confidence) scores = np.array(scores) detections = np.array(detections) features = deep_sort_encode(im, detections.copy()) detections = [ Detection(bbox, score, feature) for bbox, score, feature in zip(detections, scores, features) ] # Run non-maxima suppression. boxes = np.array([d.tlwh for d in detections]) scores = np.array([d.confidence for d in detections]) indices = prep.non_max_suppression(boxes, nms_max_overlap, scores) detections = [detections[i] for i in indices] tracker.predict() tracker.update(detections) trackers = tracker.tracks for track in trackers: if not track.is_confirmed() or track.time_since_update > 1: continue bbox = track.to_tlbr() id_num = str(track.track_id) csv.writerow([ frame_id, id_num, int(bbox[0]), int(bbox[1]), int(bbox[2]) - int(bbox[0]), int(bbox[3]) - int(bbox[1]) ]) csv_file.flush() cv2.rectangle(im, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (0, 255, 255), thick // 3) cv2.putText(im, id_num, (int(bbox[0]), int(bbox[1]) - 12), 0, 1e-3 * h, (255, 255, 255), thick // 6) # cv2.rectangle(im,(bbox[0],bbox[1]),(bbox[2],bbox[3]),(0,0,255),2) # cv2.rectangle(im,(int(bbox[0]),int(bbox[1])-10),(int(bbox[0]+200),int(bbox[1])+10),(10,10,10),-1) # cv2.putText(im, id_num,(int(bbox[0]),int(bbox[1]-2)),cv2.FONT_HERSHEY_SIMPLEX,.45,(255,255,255))#,cv2.CV_AA) return im
def camera(self): file = self.FLAGS.demo SaveVideo = self.FLAGS.saveVideo if self.FLAGS.track: if self.FLAGS.tracker == "deep_sort": from deep_sort import generate_detections from deep_sort.deep_sort import nn_matching from deep_sort.deep_sort.tracker import Tracker metric = nn_matching.NearestNeighborDistanceMetric( "cosine", 0.2, 100) tracker = Tracker(metric) encoder = generate_detections.create_box_encoder( os.path.abspath( "deep_sort/resources/networks/mars-small128.ckpt-68577")) elif self.FLAGS.tracker == "sort": from sort.sort import Sort encoder = None tracker = Sort() if self.FLAGS.BK_MOG and self.FLAGS.track: fgbg = cv2.bgsegm.createBackgroundSubtractorMOG() if file == 'camera': file = 0 else: assert os.path.isfile(file), \ 'file {} does not exist'.format(file) camera = skvideo.io.VideoCapture(file) if file == 0: self.say('Press [ESC] to quit video') assert camera.isOpened(), \ 'Cannot capture source' if self.FLAGS.csv: f = open('{}.csv'.format(file), 'w') writer = csv.writer(f, delimiter=',') writer.writerow(['frame_id', 'track_id', 'x', 'y', 'w', 'h']) f.flush() else: f = None writer = None if file == 0: #camera window cv2.namedWindow('', 0) _, frame = camera.read() height, width, _ = frame.shape cv2.resizeWindow('', width, height) else: _, frame = camera.read() height, width, _ = frame.shape if SaveVideo: if file == 0: #camera window fps = 1 / self._get_fps(frame) if fps < 1: fps = 1 else: fps = get_fps_rate(file) output_file = 'output_{}'.format(file) if os.path.exists(output_file): os.remove(output_file) videoWriter = skvideo.io.VideoWriter(output_file, fps=fps, frameSize=(width, height)) videoWriter.open() # buffers for demo in batch buffer_inp = list() buffer_pre = list() elapsed = 0 start = timer() self.say('Press [ESC] to quit demo') #postprocessed = [] # Loop through frames n = 0 while camera.isOpened(): elapsed += 1 _, frame = camera.read() if frame is None: print('\nEnd of Video') break if self.FLAGS.skip != n: n += 1 continue n = 0 if self.FLAGS.BK_MOG and self.FLAGS.track: fgmask = fgbg.apply(frame) else: fgmask = None preprocessed = self.framework.preprocess(frame) buffer_inp.append(frame) buffer_pre.append(preprocessed) # Only process and imshow when queue is full if elapsed % self.FLAGS.queue == 0: feed_dict = {self.inp: buffer_pre} net_out = self.sess.run(self.out, feed_dict) for img, single_out in zip(buffer_inp, net_out): if not self.FLAGS.track: postprocessed = self.framework.postprocess(single_out, img, save=False) else: postprocessed = self.framework.postprocess( single_out, img, frame_id=elapsed, csv_file=f, csv=writer, mask=fgmask, encoder=encoder, tracker=tracker, save=False) if SaveVideo: videoWriter.write(postprocessed) # Clear Buffers buffer_inp = list() buffer_pre = list() if elapsed % 5 == 0: sys.stdout.write('\r') sys.stdout.write('{0:3.3f} FPS'.format(elapsed / (timer() - start))) sys.stdout.flush() sys.stdout.write('\n') if SaveVideo: videoWriter.release() if self.FLAGS.csv: f.close() camera.release()
def run(sequence_dir, detection_file, output_file, min_confidence, nms_max_overlap, min_detection_height, max_cosine_distance, nn_budget, display): """Run multi-target tracker on a particular sequence. Parameters ---------- sequence_dir : str Path to the MOTChallenge sequence directory. detection_file : str Path to the detections file. output_file : str Path to the tracking output file. This file will contain the tracking results on completion. min_confidence : float Detection confidence threshold. Disregard all detections that have a confidence lower than this value. nms_max_overlap: float Maximum detection overlap (non-maxima suppression threshold). min_detection_height : int Detection height threshold. Disregard all detections that have a height lower than this value. max_cosine_distance : float Gating threshold for cosine distance metric (object appearance). nn_budget : Optional[int] Maximum size of the appearance descriptor gallery. If None, no budget is enforced. display : bool If True, show visualization of intermediate tracking results. """ seq_info = gather_sequence_info(sequence_dir, detection_file) metric = nn_matching.NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget) tracker = Tracker(metric) results = [] def frame_callback(vis, frame_idx): print("Processing frame %05d" % frame_idx) # Load image and generate detections. detections = create_detections(seq_info["detections"], frame_idx, min_detection_height) detections = [d for d in detections if d.confidence >= min_confidence] # Run non-maxima suppression. boxes = np.array([d.tlwh for d in detections]) scores = np.array([d.confidence for d in detections]) indices = preprocessing.non_max_suppression(boxes, nms_max_overlap, scores) detections = [detections[i] for i in indices] # Update tracker. tracker.predict() tracker.update(detections) # Update visualization. if display: image = cv2.imread(seq_info["image_filenames"][frame_idx], cv2.IMREAD_COLOR) vis.set_image(image.copy()) vis.draw_detections(detections) vis.draw_trackers(tracker.tracks) # Store results. for track in tracker.tracks: if not track.is_confirmed() or track.time_since_update > 1: continue bbox = track.to_tlwh() results.append([ frame_idx, track.track_id, bbox[0], bbox[1], bbox[2], bbox[3] ]) # Run tracker. if display: visualizer = visualization.Visualization(seq_info, update_ms=5) else: visualizer = visualization.NoVisualization(seq_info) visualizer.run(frame_callback) # Store results. f = open(output_file, 'w') for row in results: print('%d,%d,%.2f,%.2f,%.2f,%.2f,1,-1,-1,-1' % (row[0], row[1], row[2], row[3], row[4], row[5]), file=f)
def __init__(self, max_cosine_distance=0.2, nn_budget=100): metric = nn_matching.NearestNeighborDistanceMetric( "cosine", max_cosine_distance, nn_budget) self.tracker = Tracker(metric) self.results = []
'tracker': 'deep_sort', 'skip': 0, 'csv': False } buff = deque(maxlen=20) tfnet = TFNet(FLAGS) # from piyush import Connect file = FLAGS["demo"] SaveVideo = FLAGS["saveVideo"] # w = Connect() if FLAGS["track"]: if FLAGS["tracker"] == "deep_sort": from deep_sort import generate_detections from deep_sort.deep_sort import nn_matching from deep_sort.deep_sort.tracker import Tracker metric = nn_matching.NearestNeighborDistanceMetric("cosine", 0.2, 100) tracker = Tracker(metric) encoder = generate_detections.create_box_encoder( os.path.abspath( "deep_sort/resources/networks/mars-small128.ckpt-68577")) elif FLAGS["tracker"] == "sort": from sort.sort import Sort encoder = None tracker = Sort() if FLAGS["BK_MOG"] and FLAGS["track"]: fgbg = cv2.bgsegm.createBackgroundSubtractorMOG() if file == 'camera': file = 0 else: assert os.path.isfile(file), \
def camera(self): file = self.FLAGS.demo SaveVideo = self.FLAGS.saveVideo if self.FLAGS.track: if self.FLAGS.tracker == "deep_sort": from deep_sort import generate_detections from deep_sort.deep_sort import nn_matching from deep_sort.deep_sort.tracker import Tracker metric = nn_matching.NearestNeighborDistanceMetric( "cosine", 0.2, 100) tracker = Tracker(metric) encoder = generate_detections.create_box_encoder( os.path.abspath( "deep_sort/resources/networks/mars-small128.ckpt-68577")) elif self.FLAGS.tracker == "sort": from sort.sort import Sort encoder = None tracker = Sort() if self.FLAGS.BK_MOG and self.FLAGS.track: fgbg = cv2.bgsegm.createBackgroundSubtractorMOG() if file == 'camera': file = 0 else: assert os.path.isfile(file), \ 'file {} does not exist'.format(file) camera = cv2.VideoCapture(file) if file == 0: self.say('Press [ESC] to quit video') assert camera.isOpened(), \ 'Cannot capture source' if self.FLAGS.csv: f = open('{}.csv'.format(file), 'w') writer = csv.writer(f, delimiter=',') writer.writerow(['frame_id', 'track_id', 'x', 'y', 'w', 'h']) f.flush() else: f = None writer = None if file == 0: #camera window cv2.namedWindow(self.FLAGS.object_id, 0) _, frame = camera.read() height, width, _ = frame.shape cv2.resizeWindow(self.FLAGS.object_id, width * 0.5, height * 0.5) else: _, frame = camera.read() height, width, _ = frame.shape if self.FLAGS.push_stream: ffmpeg_pipe(self, file, width, height) if SaveVideo: fourcc = cv2.VideoWriter_fourcc(*'XVID') if file == 0: #camera window fps = 1 / self._get_fps(frame) if fps < 1: fps = 1 else: fps = round(camera.get(cv2.CAP_PROP_FPS)) videoWriter = cv2.VideoWriter('output_{}'.format(file), fourcc, fps, (width, height)) # buffers for demo in batch buffer_inp = list() buffer_pre = list() elapsed = 0 start = timer() self.say('Press [ESC] to quit demo') #postprocessed = [] # Loop through frames n = 0 while camera.isOpened(): if self.FLAGS.process_status == 1: # esc print("gongjia: Stoped! ") break if self.FLAGS.process_status == 2: #print("gongjia: Paused! ") continue elapsed += 1 _, frame = camera.read() if frame is None: print('\nEnd of Video') break if self.FLAGS.skip != n: n += 1 continue n = 0 if self.FLAGS.BK_MOG and self.FLAGS.track: fgmask = fgbg.apply(frame) else: fgmask = None preprocessed = self.framework.preprocess(frame) buffer_inp.append(frame) buffer_pre.append(preprocessed) # Only process and imshow when queue is full if elapsed % self.FLAGS.queue == 0: feed_dict = {self.inp: buffer_pre} net_out = self.sess.run(self.out, feed_dict) for img, single_out in zip(buffer_inp, net_out): if not self.FLAGS.track: postprocessed = self.framework.postprocess(single_out, img) else: postprocessed = self.framework.postprocess( single_out, img, frame_id=elapsed, csv_file=f, csv=writer, mask=fgmask, encoder=encoder, tracker=tracker) if SaveVideo: videoWriter.write(postprocessed) if self.FLAGS.display: cv2.imshow(self.FLAGS.object_id, postprocessed) if self.FLAGS.push_stream: #im = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) self.pipe.stdin.write(postprocessed.tobytes()) # Clear Buffers buffer_inp = list() buffer_pre = list() if elapsed % 5 == 0: sys.stdout.write('\r') sys.stdout.write(' {0:3.3f} FPS '.format(elapsed / (timer() - start))) sys.stdout.flush() if self.FLAGS.display: choice = cv2.waitKey(1) if choice == 27: break cv2.imwrite( '{}_{}_counter.jpg'.format(self.FLAGS.demo, self.FLAGS.object_id), postprocessed) sys.stdout.write('\n') if SaveVideo: videoWriter.release() if self.FLAGS.csv: f.close() camera.release() if self.FLAGS.display: cv2.destroyAllWindows() if self.FLAGS.push_stream: self.pipe.stdin.close() self.pipe.wait()
def camera(self): file = self.FLAGS.demo SaveVideo = self.FLAGS.saveVideo if self.FLAGS.track : if self.FLAGS.tracker == "deep_sort": from deep_sort import generate_detections from deep_sort.deep_sort import nn_matching from deep_sort.deep_sort.tracker import Tracker metric = nn_matching.NearestNeighborDistanceMetric( "cosine", 0.2, 100) tracker = Tracker(metric) encoder = generate_detections.create_box_encoder( os.path.abspath("deep_sort/resources/networks/mars-small128.ckpt-68577")) elif self.FLAGS.tracker == "sort": from sort.sort import Sort encoder = None tracker = Sort() if self.FLAGS.BK_MOG and self.FLAGS.track : fgbg = cv2.bgsegm.createBackgroundSubtractorMOG() camera = cv2.VideoCapture(file[0]) camera1 = cv2.VideoCapture(file[1]) if file == 0: self.say('Press [ESC] to quit video') assert camera.isOpened(), \ 'Cannot capture source' if self.FLAGS.csv : f = open('{}.csv'.format(file),'w') writer = csv.writer(f, delimiter=',') writer.writerow(['frame_id', 'track_id' , 'x', 'y', 'w', 'h']) f.flush() else : f =None writer= None if file == 0:#camera window cv2.namedWindow('', 0) _, frame = camera.read() height, width, _ = frame.shape cv2.resizeWindow('', width, height) else: ret, frame = camera.read() ret1, frame1 = camera1.read() height, width, _ = (frame.shape[0], (frame.shape[1]+frame1.shape[1]), 3) if SaveVideo: fourcc = cv2.VideoWriter_fourcc(*'XVID') if file == 0:#camera window fps = 1 / self._get_fps(frame) if fps < 1: fps = 1 else: fps = round(camera.get(cv2.CAP_PROP_FPS)) videoWriter = cv2.VideoWriter( "/".join(filepath.split("/")[:-1]) + '/output_{}'.format(filepath.split("/")[-1]), fourcc, fps, (width, height)) # buffers for demo in batch buffer_inp = list() buffer_pre = list() elapsed = 0 start = timer() self.say('Press [ESC] to quit demo') #postprocessed = [] # Loop through frames n = 0 while (camera.isOpened() and camera1.isOpened()): elapsed += 1 _, frame = camera.read() _, frame1 = camera1.read() if frame is None: print ('\nEnd of Video') break if self.FLAGS.skip != n : n+=1 continue n = 0 if self.FLAGS.BK_MOG and self.FLAGS.track : fgmask = fgbg.apply(frame) else : fgmask = None vis = np.concatenate((frame, frame1), axis=1) preprocessed = self.framework.preprocess(vis) buffer_inp.append(vis) buffer_pre.append(preprocessed) # Only process and imshow when queue is full if elapsed % self.FLAGS.queue == 0: feed_dict = {self.inp: buffer_pre} net_out = self.sess.run(self.out, feed_dict) for img, single_out in zip(buffer_inp, net_out): if not self.FLAGS.track : postprocessed = self.framework.postprocess( single_out, img) else : postprocessed = self.framework.postprocess( single_out, img, frame_id = elapsed, csv_file=f,csv=writer,mask = fgmask, encoder=encoder,tracker=tracker) # postprocessed1 = self.framework.postprocess( # single_out, img[:height,int(width/2):width, :3],frame_id = elapsed, # csv_file=f,csv=writer,mask = fgmask, # encoder=encoder,tracker=tracker) if SaveVideo: videoWriter.write(postprocessed) if self.FLAGS.display : cv2.imshow('', postprocessed) # Clear Buffers buffer_inp = list() buffer_pre = list() if elapsed % 5 == 0: sys.stdout.write('\r') sys.stdout.write('{0:3.3f} FPS'.format( elapsed / (timer() - start))) sys.stdout.flush() if self.FLAGS.display : choice = cv2.waitKey(1) if choice == 27: break sys.stdout.write('\n') if SaveVideo: videoWriter.release() if self.FLAGS.csv : f.close() camera.release() if self.FLAGS.display : cv2.destroyAllWindows()
def detect_video(graph, yolo, encoder, video_filepath, mark_on_video, show): global_object_ids = dict() object_id_counter = 0 # Definition of the parameters max_cosine_distance = 0.3 nn_budget = None nms_max_overlap = 0.3 # 1.0 # read video filepath cap = cv2.VideoCapture(video_filepath) try: # per class object tracker class_tracker_dict = dict() i_frame = -1 detection_dict = {} while True: ret, frame = cap.read() i_frame += 1 if not ret: break with graph.as_default(): out_boxes, out_scores, out_classes = detect_frame( yolo=yolo, image=Image.fromarray(frame)) # convert to [x,y,w,h] boxs = np.array([[bb[1], bb[0], bb[3] - bb[1], bb[2] - bb[0]] for bb in out_boxes]) # add detections to class tracker dict for class_id in np.unique(out_classes): if class_id not in class_tracker_dict: metric = nn_matching.NearestNeighborDistanceMetric( "cosine", max_cosine_distance, nn_budget) class_tracker_dict[class_id] = Tracker(metric) # update all trackers with incoming detections for class_id, tracker in class_tracker_dict.items(): inds = out_classes == class_id with graph.as_default(): features = encoder(frame, boxs[inds]) detections = [ Detection(bbox, score, feature) for bbox, score, feature in zip( boxs[inds], out_scores[inds], features) ] # Run non-maxima suppression. boxes = np.array([d.tlwh for d in detections]) scores = np.array([d.confidence for d in detections]) indices = preprocessing.non_max_suppression( boxes, nms_max_overlap, scores) detections = [detections[i] for i in indices] # Call the tracker class_tracker_dict[class_id].predict() class_tracker_dict[class_id].update(detections) # Save detection per frame per object_id detection_dict[i_frame] = list() for class_id, tracker in class_tracker_dict.items(): for track in tracker.tracks: # if not track.is_confirmed() or track.time_since_update > 1: if track.time_since_update > 2: continue bbox = track.to_tlbr() object_name = '{}_{}'.format(yolo.class_names[class_id], track.track_id) if object_name not in global_object_ids: global_object_ids[object_name] = object_id_counter object_id_counter += 1 # mark in video label = yolo.class_names[class_id] left = int(bbox[0]) top = int(bbox[1]) right = int(bbox[2]) bottom = int(bbox[3]) object_id = track.track_id if mark_on_video: cv2.rectangle(frame, (left, top), (right, bottom), (255, 255, 255), 2) cv2.putText(frame, '{}_{}'.format(label, object_id), (int(bbox[0]), int(bbox[1])), 0, 1, (0, 255, 0), 1) # create results dictionary detection_dict[i_frame].append({ 'top': top, 'left': left, 'right': right, 'bottom': bottom, 'label': label, 'object_id': global_object_ids[object_name] }) if show: cv2.namedWindow("result", cv2.WINDOW_AUTOSIZE) cv2.imshow("result", frame) if cv2.waitKey(1) & 0xFF == ord('q'): break finally: cap.release() return detection_dict
def camera(self): file = self.FLAGS.demo SaveVideo = self.FLAGS.saveVideo if self.FLAGS.track: if self.FLAGS.tracker == "deep_sort": from deep_sort import generate_detections from deep_sort.deep_sort import nn_matching from deep_sort.deep_sort.tracker import Tracker metric = nn_matching.NearestNeighborDistanceMetric( "cosine", 0.2, 100) tracker = Tracker(metric) encoder = generate_detections.create_box_encoder( os.path.abspath( "deep_sort/resources/networks/mars-small128.ckpt-68577")) elif self.FLAGS.tracker == "sort": from sort.sort import Sort encoder = None tracker = Sort() if self.FLAGS.BK_MOG and self.FLAGS.track: fgbg = cv2.bgsegm.createBackgroundSubtractorMOG() # if file == 'camera': # file = 0 # else: # assert os.path.isfile(file), \ # 'file {} does not exist'.format(file) camera1 = cv2.VideoCapture(file[0]) camera2 = cv2.VideoCapture(file[1]) camera3 = cv2.VideoCapture(file[2]) # if file == 0: # self.say('Press [ESC] to quit video') # # assert camera.isOpened(), \ # 'Cannot capture source' if self.FLAGS.csv: f = open('{}.csv'.format(file), 'w') writer = csv.writer(f, delimiter=',') writer.writerow(['frame_id', 'track_id', 'x', 'y', 'w', 'h']) f.flush() else: f = None writer = None # if file == 0:#camera window # cv2.namedWindow('', 0) # _, frame = camera.read() # height, width, _ = frame.shape # cv2.resizeWindow('', width, height) # else: # _, frame = camera.read() # height, width, _ = frame.shape # if SaveVideo: # fourcc = cv2.VideoWriter_fourcc(*'XVID') # if file == 0:#camera window # fps = 1 / self._get_fps(frame) # if fps < 1: # fps = 1 # else: # fps = round(camera1.get(cv2.CAP_PROP_FPS)) # videoWriter = cv2.VideoWriter( # 'output_{}'.format(file), fourcc, fps, (width, height)) # buffers for demo in batch buffer_inp = list() buffer_pre = list() elapsed = 0 start = timer() self.say('Press [ESC] to quit demo') #postprocessed = [] # Loop through frames n = 0 while (camera1.isOpened() and camera2.isOpened() and camera3.isOpened()): elapsed += 1 ret1, frame1 = camera1.read() ret2, frame2 = camera2.read() ret3, frame3 = camera3.read() #if(ret1 and ret2 and ret3): h1, w1 = frame1.shape[:2] vis = np.concatenate((frame2, frame1, frame3), axis=1) if self.FLAGS.skip != n: n += 1 continue n = 0 # while camera.isOpened(): # elapsed += 1 # _, frame = camera.read() # if frame is None: # print ('\nEnd of Video') # break # if self.FLAGS.skip != n : # n+=1 # continue # n = 0 if self.FLAGS.BK_MOG and self.FLAGS.track: fgmask = fgbg.apply(vis) else: fgmask = None preprocessed = self.framework.preprocess(vis) buffer_inp.append(vis) buffer_pre.append(preprocessed) # Only process and imshow when queue is full if elapsed % self.FLAGS.queue == 0: feed_dict = {self.inp: buffer_pre} net_out = self.sess.run(self.out, feed_dict) for img, single_out in zip(buffer_inp, net_out): if not self.FLAGS.track: postprocessed = self.framework.postprocess(single_out, img) else: #print("else hi") postprocessed = self.framework.postprocess( single_out, img, frame_id=elapsed, csv_file=f, csv=writer, mask=fgmask, encoder=encoder, tracker=tracker) if SaveVideo: videoWriter.write(postprocessed) if self.FLAGS.display: cv2.imshow('This is postprocessed', postprocessed) # Clear Buffers buffer_inp = list() buffer_pre = list() if elapsed % 5 == 0: sys.stdout.write('\r') sys.stdout.write('{0:3.3f} FPS'.format(elapsed / (timer() - start))) sys.stdout.flush() if self.FLAGS.display: choice = cv2.waitKey(1) if choice == 27: break sys.stdout.write('\n') if SaveVideo: videoWriter.release() if self.FLAGS.csv: f.close() camera1.release() camera2.release() camera3.release() if self.FLAGS.display: cv2.destroyAllWindows()
def camera(self): file = self.FLAGS.demo SaveVideo = self.FLAGS.saveVideo detectedObjects = [] if self.FLAGS.track : if self.FLAGS.tracker == "deep_sort": from deep_sort import generate_detections from deep_sort.deep_sort import nn_matching from deep_sort.deep_sort.tracker import Tracker metric = nn_matching.NearestNeighborDistanceMetric( "cosine", 0.2, 100) tracker = Tracker(metric) encoder = generate_detections.create_box_encoder( os.path.abspath("deep_sort/resources/networks/mars-small128.ckpt-68577")) elif self.FLAGS.tracker == "sort": from sort.sort import Sort encoder = None tracker = Sort() if self.FLAGS.BK_MOG and self.FLAGS.track : fgbg = cv2.bgsegm.createBackgroundSubtractorMOG() if file == 'camera': file = 0 else: assert os.path.isfile(file), \ 'file {} does not exist'.format(file) camera = cv2.VideoCapture(file) # if file == 0: # self.say('Press [ESC] to quit video') assert camera.isOpened(), \ 'Cannot capture source' savedPath = 'result/output_{}'.format(os.path.basename(file)) if self.FLAGS.csv : f = open('{}.csv'.format(file),'w') writer = csv.writer(f, delimiter=',') writer.writerow(['frame_id', 'track_id' , 'x', 'y', 'w', 'h']) f.flush() else : f =None writer= None if file == 0:#camera window cv2.namedWindow('', 0) _, frame = camera.read() height, width, _ = frame.shape cv2.resizeWindow('', width, height) else: _, frame = camera.read() height, width, _ = frame.shape if SaveVideo: fourcc = cv2.VideoWriter_fourcc(*'XVID') if file == 0:#camera window fps = 1 / self._get_fps(frame) if fps < 1: fps = 1 else: fps = round(camera.get(cv2.CAP_PROP_FPS)) videoWriter = cv2.VideoWriter( savedPath, fourcc, fps, (width, height)) # buffers for demo in batch buffer_inp = list() buffer_pre = list() elapsed = 0 start = timer() # self.say('Press [ESC] to quit demo') #postprocessed = [] # Loop through frames n = 0 while camera.isOpened(): elapsed += 1 _, frame = camera.read() if frame is None: # print ('\nEnd of Video') break if self.FLAGS.skip != n : n+=1 continue n = 0 if self.FLAGS.BK_MOG and self.FLAGS.track : fgmask = fgbg.apply(frame) else : fgmask = None preprocessed = self.framework.preprocess(frame) buffer_inp.append(frame) buffer_pre.append(preprocessed) # Only process and imshow when queue is full if elapsed % self.FLAGS.queue == 0: feed_dict = {self.inp: buffer_pre} net_out = self.sess.run(self.out, feed_dict) for img, single_out in zip(buffer_inp, net_out): if not self.FLAGS.track : postprocessed = self.framework.postprocess( single_out, img) else : postprocessedTuple = self.framework.postprocess( single_out, img,frame_id = elapsed, csv_file=f,csv=writer,mask = fgmask, encoder=encoder,tracker=tracker) postprocessed = postprocessedTuple[0] detectedObjects.append(postprocessedTuple[1]) if SaveVideo: videoWriter.write(postprocessed) if self.FLAGS.display : cv2.imshow('', postprocessed) # Clear Buffers buffer_inp = list() buffer_pre = list() if elapsed % 5 == 0: # sys.stdout.write('\r') # sys.stdout.write('{0:3.3f} FPS'.format( # elapsed / (timer() - start))) sys.stdout.flush() if self.FLAGS.display : choice = cv2.waitKey(1) if choice == 27: break flattenObjects = sum(detectedObjects, []) frameDictionary = dict() numDictionary = dict() totalFrame = int(camera.get(7)) #CAP_PROP_FRAME_COUNT fps = int(camera.get(5)) #CAP_PROP_FPS illegalSecond = 30 stopFrameCountThreshold = fps * illegalSecond for object in flattenObjects: frameDictionary\ .setdefault(object.frame, []).append(object.num) numDictionary\ .setdefault(object.num, Car(object.num, Segment(object.frame, object.position), stopFrameCountThreshold))\ .update(Segment(object.frame, object.position)) resultJSON = { "frames": list(map(lambda key: {"id": key, "carNums": frameDictionary[key]}, frameDictionary)), "cars": list(map(lambda value: json.loads(json.dumps(value, default=lambda o: o.__dict__)), numDictionary.values())), "resultVideoPath": savedPath } print(json.dumps(resultJSON, indent=2)) if SaveVideo: videoWriter.release() if self.FLAGS.csv : f.close() camera.release() if self.FLAGS.display : cv2.destroyAllWindows()
def efficientDet_video_inference(video_src,compound_coef = 0,force_input_size=None, frame_skipping = 3, threshold=0.2,out_path=None,imshow=False, display_fps=False): #deep-sort variables # Definition of the parameters max_cosine_distance = 0.3 nn_budget = None nms_max_overlap = 1.0 model_filename = '/home/shaheryar/Desktop/Projects/Football-Monitoring/deep_sort/model_weights/mars-small128.pb' encoder = gdet.create_box_encoder(model_filename, batch_size=1) metric = nn_matching.NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget) tracker = Tracker(metric,n_init=5) # efficientDet-pytorch variables iou_threshold = 0.4 use_cuda = True use_float16 = False cudnn.fastest = True cudnn.benchmark = True input_size = input_sizes[compound_coef] if force_input_size is None else force_input_size # load model model = EfficientDetBackbone(compound_coef=compound_coef, num_classes=len(obj_list)) model.load_state_dict(torch.load(f'weights/efficientdet-d{compound_coef}.pth')) model.requires_grad_(False) model.eval() if use_cuda: model = model.cuda() if use_float16: model = model.half() regressBoxes = BBoxTransform() clipBoxes = ClipBoxes() # Video capture cap = cv2.VideoCapture(video_src) frame_width = int(cap.get(3)) frame_height = int(cap.get(4)) fourcc = cv2.VideoWriter_fourcc(*'MPEG') fps = cap.get(cv2.CAP_PROP_FPS) print("Video fps",fps) if(out_path is not None): outp = cv2.VideoWriter(out_path, fourcc, fps, (frame_width, frame_height)) i=0 start= time.time() current_frame_fps=0 while True: ret, frame = cap.read() if not ret: break t1=time.time() if (frame_skipping==0 or i%frame_skipping==0): # if(True): # frame preprocessing (running detections) ori_imgs, framed_imgs, framed_metas, t1 = preprocess_video(frame, width=input_size, height=input_size) if use_cuda: x = torch.stack([fi.cuda() for fi in framed_imgs], 0) else: x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0) # model predict t1=time.time() with torch.no_grad(): features, regression, classification, anchors = model(x) out = postprocess(x, anchors, regression, classification, regressBoxes, clipBoxes, threshold, iou_threshold) # Post processing out = invert_affine(framed_metas, out) # decoding bbox ,object name and scores boxes,classes,scores =decode_predictions(out[0]) org_boxes = boxes.copy() t2 = time.time() - t1 # feature extraction for deep sort boxes = [convert_bbox_to_deep_sort_format(frame.shape, b) for b in boxes] features = encoder(frame,boxes) detections = [Detection(bbox, 1.0, feature) for bbox, feature in zip(boxes, features)] boxes = np.array([d.tlwh for d in detections]) # print(boxes) scores = np.array([d.confidence for d in detections]) indices = preprocessing.non_max_suppression(boxes, nms_max_overlap, scores) detections = [detections[i] for i in indices] tracker.predict() tracker.update(detections) i = i + 1 img_show=frame.copy() for j in range(len(org_boxes)): img_show =drawBoxes(img_show,org_boxes[j],(255,255,0),str(tracker.tracks[j].track_id)) for track in tracker.tracks: if not track.is_confirmed() or track.time_since_update > 1: continue bbox = track.to_tlbr() x1=int(bbox[0]) y1 = int(bbox[1]) x2 = int(bbox[2]) y2=int(bbox[3]) roi= frame[y1:y2,x1:x2] cv2.rectangle(img_show, (x1, y1), (x2, y2), update_color_association(roi, track.track_id), 2) cv2.putText(img_show, str(track.track_id), (x1, y1), 0, 5e-3 * 100, (255, 255, 0), 1) if display_fps: current_frame_fps=1/t2 else: current_frame_fps=0 cv2.putText(img_show, 'FPS: {0:.2f}'.format(current_frame_fps), (30, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 0), 2, cv2.LINE_AA) if (i % int(fps) == 0): print("Processed ", str(int(i / fps)), "seconds") print("Time taken",time.time()-start) # print(color_dict) if imshow: img_show=cv2.resize(img_show,(0,0),fx=0.75,fy=0.75) cv2.imshow('Frame',img_show) # Press Q on keyboard to exit if cv2.waitKey(1) & 0xFF == ord('q'): break if out_path is not None: outp.write(img_show) cap.release() outp.release()
def main(): model_file = 'mars-small128.pb' score_threshold = 0 nn_budget = None max_cosine_distance = 0.2 min_confidence = 0.8 nms_max_overlap = 1.0 min_detection_height = 0 metric = nn_matching.NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget) tracker = Tracker(metric) DEFAULT_UPDATE_MS = 20 consumer = KafkaConsumer( topic_in, bootstrap_servers=['kafka1:19091'] ) # Start up producer producer = KafkaProducer(bootstrap_servers='kafka1:19091') results = [] encoder = create_box_encoder(model_file, batch_size=32) image_id = -1 for msg in consumer: buf = base64.b64decode(msg.value) payload = pickle.loads(buf) image_id, img, predictions = payload['image_id'], payload['img'], np.array(payload['frame_results']) img = from_base64(img) image = Image.fromarray(img) image = np.array(image) if len(predictions)>0: predictions = predictions[predictions[:, 6] > score_threshold, :] detections = generate_detections(encoder, image, predictions) else: detections = None results += run( image_id, image, detections, metric, tracker, min_confidence, nms_max_overlap, min_detection_height, max_cosine_distance, nn_budget) seq_info = { "min_frame_idx": 0, "max_frame_idx": np.inf, "update_ms": DEFAULT_UPDATE_MS, "image_filenames": 'video_stream', "sequence_name": 'video_stream', "image_size": img.shape, "feature_dim": detections.shape[1] - 10 if detections is not None else 0, } # # Store results. # f = open(output_file, 'w') # for row in results: # print('%d,%d,%.2f,%.2f,%.2f,%.2f,1,-1,-1,-1' % ( # row[0], row[1], row[2], row[3], row[4], row[5]),file=f) visualizer = visualization.StreamVisualization(seq_info, seq_info["update_ms"]) results = np.array(results) img_output = visualizer.run(frame_annotation_callback, image_id, image, detections, results) results = results.tolist() # Convert image to jpg _, buffer = cv2.imencode('.jpg', img_output) producer.send(topic_out, base64.b64encode(buffer))