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()