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, 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 __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): 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"))
class track_deepsort(object): 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 return_tracking_id(self, bboxes, frame): detections = [] scores = [] new_bboxes = [] nms_max_overlap = 0.1 if type(frame) is not np.ndarray: imgcv = cv2.imread(frame) else: imgcv = frame h, w, _ = frame.shape thick = int((h + w) // 100) for b in bboxes[0]: left, top, right, bot, confidence, _, _ = b detections.append( np.array([left, top, right - left, bot - top]).astype(np.float64)) scores.append(confidence) detections = np.array(detections) if detections.shape[0] == 0: return frame scores = np.array(scores) features = self.encoder(imgcv, 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] self.tracker.predict() self.tracker.update(detections) trackers = self.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) new_bboxes.append( (id_num, int(bbox[0]), int(bbox[1]), (int(bbox[2])), int(bbox[3]), track.confidence)) print(new_bboxes) return new_bboxes
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 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 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, 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, 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='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]), ])
class deepsort_rbc(): 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 reset_tracker(self): self.tracker = Tracker(self.metric) #Older yolo used to give y,x,h,w. """ def format_yolo_output( self,out_boxes): out_boxes=np.array([out_boxes[:,1],out_boxes[:,0],\ out_boxes[:,3]-out_boxes[:,1],out_boxes[:,2]-out_boxes[:,0]]) out_boxes=out_boxes.T return out_boxes """ #Current Yolo gives x_center,y_center,w,h #Deep sort needs the format `top_left_x, top_left_y, width,height def format_yolo_output(self, out_boxes): for b in range(len(out_boxes)): out_boxes[b][0] = out_boxes[b][0] - out_boxes[b][2] / 2 out_boxes[b][1] = out_boxes[b][1] - out_boxes[b][3] / 2 return out_boxes def run_deep_sort(self, frame, out_scores, out_boxes, out_classes): out_boxes = self.format_yolo_output(out_boxes) if out_boxes == []: self.tracker.predict() trackers = self.tracker.tracks return trackers detections = np.array(out_boxes) #features = self.encoder(frame, detections.copy()) features = self.encoder.extract_features(frame, detections) #print(frame.shape) detections = [Detection(bbox, score, feature,classname) \ for bbox,score, feature,classname in\ zip(detections,out_scores, features,out_classes)] outboxes = np.array([d.tlwh for d in detections]) outscores = np.array([d.confidence for d in detections]) indices = prep.non_max_suppression(outboxes, 0.8, outscores) detections = [detections[i] for i in indices] self.tracker.predict() self.tracker.update(detections) trackers = self.tracker.tracks return trackers
class deepsort_rbc(): def __init__(self): self.metric = nn_matching.NearestNeighborDistanceMetric( "cka", .5, 30) #metric, matching_threshold, budget self.tracker = Tracker(self.metric) def reset_tracker(self): self.tracker = Tracker(self.metric) #Deep sort needs the format `top_left_x, top_left_y, width,height def format_boxes(self, out_boxes): for b in range(len(out_boxes)): out_boxes[b][2] = out_boxes[b][2] - out_boxes[b][0] out_boxes[b][3] = out_boxes[b][3] - out_boxes[b][1] return out_boxes def run_deep_sort(self, frame, out_boxes, out_scores, classIDs, scales, ids, features): if out_boxes == []: self.tracker.predict() print('No detections') trackers = self.tracker.tracks return trackers wh = np.flip(frame.shape[0:2]) out_boxes[:, 0:2] = (out_boxes[:, 0:2] * wh).astype(float) out_boxes[:, 2:4] = (out_boxes[:, 2:4] * wh).astype(float) # Give format boxes out_boxes = self.format_boxes(out_boxes) # Create a detection object to parse in deepsort dets = [Detection(bbox, score, feature) \ for bbox, score, feature in\ zip(out_boxes, out_scores, features)] outboxes = np.array([d.tlwh for d in dets]) outscores = np.array([d.confidence for d in dets]) # Non max suppression including confidence # works best using pixel coordinates, 0.3<=overlap<=0.5 indices = non_max_suppression(outboxes, 0.4, outscores) dets = [dets[i] for i in indices] classIDs = [classIDs[i] for i in indices] scales = [scales[i] for i in indices] ids = [ids[i] for i in indices] detections_class = (dets, classIDs, scales, ids) # DeepSort cycle self.tracker.predict() self.tracker.update(dets) return self.tracker, detections_class
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 __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 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))
class DeepSortNode(Observer, Subject): objectBoundingBoxes = [] objectIds = [] image = [] END = False 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 update(self, subject): if (subject.END): self.END = True self.notify() return # update tracker with detection from detector boundBoxes = subject.rois confidences = subject.scores if (self.applyMask): features = self.encoder(subject.image, np.array(boundBoxes), subject.masks) else: features = self.encoder(subject.image, np.array(boundBoxes)) detections = [ Detection(bbox, confidence, feature) for bbox, confidence, feature in zip( boundBoxes, confidences, features) ] self.tracker.predict() self.tracker.update(detections) # extract bounding boxes and Ids self.image = subject.image self.objectBoundingBoxes = [] self.objectIds = [] tracks = self.tracker.tracks for track in tracks: if not track.is_confirmed() or track.time_since_update > 1: continue self.objectBoundingBoxes.append(track.to_tlbr()) self.objectIds.append(str(track.track_id)) # notify deep sort event listeners self.notify() def notify(self): print("DeepSortNode: update people tracking") for observer in self.observers: observer.update(self)
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 __init__(self): self.metric = nn_matching.NearestNeighborDistanceMetric( "cka", .5, 30) #metric, matching_threshold, budget self.tracker = Tracker(self.metric)
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 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()
class Deepsort_original(object): 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 reset_tracker(self): self.tracker = Tracker(self.metric) # Deep sort needs the format `top_left_x, top_left_y, width,height def format_yolo_output(self, out_boxes): for b in range(len(out_boxes)): out_boxes[b][0] = out_boxes[b][0] - out_boxes[b][2] / 2 out_boxes[b][1] = out_boxes[b][1] - out_boxes[b][3] / 2 return out_boxes def pre_process(self, frame, detections): transforms = torchvision.transforms.Compose([ torchvision.transforms.ToPILImage(), torchvision.transforms.Resize((128, 64)), torchvision.transforms.ToTensor(), ]) crops = [] for d in detections: for i in range(len(d)): if d[i] < 0: d[i] = 0 img_h, img_w, img_ch = frame.shape xmin, ymin, w, h = d if xmin > img_w: xmin = img_w if ymin > img_h: ymin = img_h xmax = xmin + w ymax = ymin + h ymin = abs(int(ymin)) ymax = abs(int(ymax)) xmin = abs(int(xmin)) xmax = abs(int(xmax)) try: crop = frame[ymin:ymax, xmin:xmax, :] crop = transforms(crop) crops.append(crop) except: continue crops = torch.stack(crops) return crops def extract_features_only(self, frame, coords): for i in range(len(coords)): if coords[i] < 0: coords[i] = 0 img_h, img_w, img_ch = frame.shape xmin, ymin, w, h = coords if xmin > img_w: xmin = img_w if ymin > img_h: ymin = img_h xmax = xmin + w ymax = ymin + h ymin = abs(int(ymin)) ymax = abs(int(ymax)) xmin = abs(int(xmin)) xmax = abs(int(xmax)) crop = frame[ymin:ymax, xmin:xmax, :] # crop = crop.astype(np.uint8) # print(crop.shape,[xmin,ymin,xmax,ymax],frame.shape) crop = self.transforms(crop) crop = crop.cuda() gaussian_mask = self.gaussian_mask input_ = crop * gaussian_mask input_ = torch.unsqueeze(input_, 0) features = self.encoder.forward_once(input_) features = features.detach().cpu().numpy() corrected_crop = [xmin, ymin, xmax, ymax] return features, corrected_crop def run_deep_sort(self, frame, out_scores, out_boxes): if out_boxes == []: self.tracker.predict() print('No detections') trackers = self.tracker.tracks return trackers detections = np.array(out_boxes) # features = self.encoder(frame, detections.copy()) processed_crops = self.pre_process(frame, detections).cuda() # processed_crops = self.gaussian_mask * processed_crops features = self.encoder.forward_once(processed_crops) features = features.detach().cpu().numpy() if len(features.shape) == 1: features = np.expand_dims(features, 0) dets = [Detection(bbox, score, feature) \ for bbox, score, feature in \ zip(detections, out_scores, features)] self.tracker.predict() self.tracker.update(dets) return self.tracker, dets
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 = []
class yact_node: 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 callback(self, msg): self.img = cv2.imdecode(np.fromstring(msg.data, np.uint8), 1) lstDetections = msg.bounding_boxes lstDetsDeepSort = [] if (lstDetections): for det in lstDetections: if det.Class == 'person': # Deep Sort Bounding Boxes # Use the format TOP LEFT WIDTH HEIGHT (tlwh) dsWidth = det.xmax - det.xmin dsHeight = det.ymax - det.ymin lstDetsDeepSort.append( [det.xmin, det.ymin, dsWidth, dsHeight]) self.detectionAndId = [] #region DEEPSORT if self.frameCount % 3 == 0: features = self.encoder(self.img, lstDetsDeepSort) # Create DeepSort detections trackedObjects = [ Detection(bbox, 1.0, feature) for bbox, feature in zip(lstDetsDeepSort, features) ] self.tracker.predict() self.tracker.update(trackedObjects) #endregion self.displayDetections() msgDetectionAndID = DetectionAndID() msgDetectionAndID.header = msg.header msgDetectionAndID.detections = self.detectionAndId #region DISPLAY if self.frameCount % 10 == 0: timer = time.time() - self.timePrev self.timePrev = time.time() self.intFPS = int(10 / timer) # Print FPS cv2.putText(self.img, 'FPS: {}'.format(self.intFPS), (self.img.shape[1] - 100, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2) cv2.imshow('yact: people_tracker.py', self.img) cv2.waitKey(3) self.frameCount += 1 #endregion #region DEBUG def displayDetections(self): for track in self.tracker.tracks: if not track.is_confirmed() or track.time_since_update > 1: continue bbox = track.to_tlbr() # Construct message msgBbid = BoundingBoxID() msgBbid.boundingBox.xmin = int(bbox[0]) msgBbid.boundingBox.ymin = int(bbox[1]) msgBbid.boundingBox.xmax = int(bbox[2]) msgBbid.boundingBox.ymax = int(bbox[3]) msgBbid.id = int(track.track_id) self.detectionAndId.append(msgBbid) # Draw bounding box and label cv2.rectangle(self.img, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (255, 0, 0), 2) labelText = 'ID: {}'.format(track.track_id) cv2.putText(self.img, labelText, (int(bbox[0]), int(bbox[1]) - 3), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
class deepsort_rbc(): 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 reset_tracker(self): self.tracker = Tracker(self.metric) #Deep sort needs the format `top_left_x, top_left_y, width,height def format_yolo_output(self, out_boxes): for b in range(len(out_boxes)): out_boxes[b][0] = out_boxes[b][0] - out_boxes[b][2] / 2 out_boxes[b][1] = out_boxes[b][1] - out_boxes[b][3] / 2 return out_boxes def pre_process(self, frame, detections): transforms = torchvision.transforms.Compose([ \ torchvision.transforms.ToPILImage(),\ torchvision.transforms.Resize((128,128)),\ torchvision.transforms.ToTensor()]) crops = [] for d in detections: for i in range(len(d)): if d[i] < 0: d[i] = 0 img_h, img_w, img_ch = frame.shape xmin, ymin, w, h = d if xmin > img_w: xmin = img_w if ymin > img_h: ymin = img_h xmax = xmin + w ymax = ymin + h ymin = abs(int(ymin)) ymax = abs(int(ymax)) xmin = abs(int(xmin)) xmax = abs(int(xmax)) try: crop = frame[ymin:ymax, xmin:xmax, :] crop = transforms(crop) crops.append(crop) except: continue crops = torch.stack(crops) return crops def extract_features_only(self, frame, coords): for i in range(len(coords)): if coords[i] < 0: coords[i] = 0 img_h, img_w, img_ch = frame.shape xmin, ymin, w, h = coords if xmin > img_w: xmin = img_w if ymin > img_h: ymin = img_h xmax = xmin + w ymax = ymin + h ymin = abs(int(ymin)) ymax = abs(int(ymax)) xmin = abs(int(xmin)) xmax = abs(int(xmax)) crop = frame[ymin:ymax, xmin:xmax, :] #crop = crop.astype(np.uint8) #print(crop.shape,[xmin,ymin,xmax,ymax],frame.shape) crop = self.transforms(crop) crop = crop.cuda() gaussian_mask = self.gaussian_mask input_ = crop * gaussian_mask input_ = torch.unsqueeze(input_, 0) features = self.encoder.forward_once(input_) features = features.detach().cpu().numpy() corrected_crop = [xmin, ymin, xmax, ymax] return features, corrected_crop def run_deep_sort(self, frame, out_scores, out_boxes): if out_boxes == []: self.tracker.predict() print('No detections') trackers = self.tracker.tracks return trackers detections = np.array(out_boxes) #features = self.encoder(frame, detections.copy()) processed_crops = self.pre_process(frame, detections).cuda() processed_crops = self.gaussian_mask * processed_crops features = self.encoder.forward_once(processed_crops) features = features.detach().cpu().numpy() if len(features.shape) == 1: features = np.expand_dims(features, 0) dets = [Detection(bbox, score, feature) \ for bbox,score, feature in\ zip(detections,out_scores, features)] outboxes = np.array([d.tlwh for d in dets]) outscores = np.array([d.confidence for d in dets]) indices = prep.non_max_suppression(outboxes, 0.8, outscores) dets = [dets[i] for i in indices] self.tracker.predict() self.tracker.update(dets) return self.tracker, dets
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 reset_tracker(self): self.tracker = Tracker(self.metric)
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()