def __init__(self,
filter_class=None,
model='yolox-s',
ckpt='wieghts/yolox_s.pth',
LineMapping=None):
self.LineMapping = LineMapping
self.detector = Predictor(model, ckpt)
cfg = get_config()
cfg.merge_from_file("deep_sort/configs/deep_sort.yaml")
self.deepsort = DeepSort(
cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST,
min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
self.filter_class = COCO_CLASSES
def load_info(config_deepsort, device, imgsz, out, weights):
# initialize deepsort
cfg = get_config()
cfg.merge_from_file(config_deepsort)
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST,
min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
# Initialize
device = select_device(device)
if not os.path.exists(out):
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = torch.load(weights,
map_location=device)['model'].float() # load to FP32
model.to(device).eval()
if half:
model.half() # to FP16
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
return t0, half, device, model, names, deepsort
def main():
video_name = 'car.mp4'
# video_name = 'car.mp4'
cap = cv2.VideoCapture(f'data/videos/{video_name}')
fource = cv2.VideoWriter_fourcc(*'mp4v')
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(f'runs/track/{video_name}.mp4', fource, 30, (width, height))
""" yolov5 目标检测器 """
yolov5_detector = YOLOv5Detector()
""" deepsort 追踪器 """
cfg = get_config()
cfg.merge_from_file("deep_sort/configs/deep_sort.yaml")
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST,
min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
window_name = 'Real-time tracking'
while True:
state, frame = cap.read()
if not state:
break
""" 检测目标 """
image, bbox_container = yolov5_detector(frame)
""" 仅保留人车信息"""
bbox_container = cut_bbox_container(bbox_container)
""" 初始化一些变量 """
xywh_bboxs = []
labels = []
confs = []
for bbox in bbox_container:
xywh_bboxs.append(xyxy_to_xywh(bbox['box']))
labels.append(bbox['class'])
confs.append(bbox['confidence'])
""" 检测到目标后才有追踪 """
if labels:
""" detections --> deepsort """
xywhs = torch.Tensor(xywh_bboxs)
confss = torch.Tensor(confs)
outputs = deepsort.update(xywhs, confss, labels, frame)
obj_ids = []
bbox_draw = []
if len(outputs) > 0:
for (x1, y1, x2, y2, label, track_id) in outputs:
bbox_draw.append({'class': label, 'box': [x1, y1, x2, y2]})
obj_ids.append(track_id)
""" 绘图显示 """
draw_image(frame, bbox_draw, obj_ids)
""" 输出一些信息 """
for info in bbox_draw:
print(info)
print(obj_ids)
print('---')
cv2.imshow(window_name, frame)
vid_writer.write(frame)
cv2.waitKey(1)
""" 点 x 退出 """
if cv2.getWindowProperty(window_name, cv2.WND_PROP_AUTOSIZE) < 1:
break
cap.release()
vid_writer.release()
cv2.destroyAllWindows()
from deep_sort.utils.parser import get_config
from deep_sort.deep_sort import DeepSort
import torch
import cv2
palette = (2**11 - 1, 2**15 - 1, 2**20 - 1)
cfg = get_config()
cfg.merge_from_file("deep_sort/configs/deep_sort.yaml")
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST,
min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
def plot_bboxes(image, bboxes, line_thickness=None):
# Plots one bounding box on image img
tl = line_thickness or round(0.002 * (image.shape[0] + image.shape[1]) /
2) + 1 # line/font thickness
for (x1, y1, x2, y2, cls_id, pos_id) in bboxes:
if cls_id in ['smoke', 'phone', 'eat']:
color = (0, 0, 255)
else:
color = (0, 255, 0)
if cls_id == 'eat':
cls_id = 'eat-drink'
c1, c2 = (x1, y1), (x2, y2)
def detect(weights='',
source='inferences/images',
output='inferences/output',
img_size=640,
conf_thres=0.4,
iou_thres=0.5,
device='',
view_img=False,
save_img=False,
save_txt=False,
classes=None,
agnostic_nms=True,
augment=True,
update=True,
fps_count=1,
line_coordinat = [], # [[(x1,y1),(x2,y2)], ...]
polygon=[], # [[[(x1,y1),(x2,y2),(x3,y3),...], [0(u)/1(r)/2(d)/3(l)]], ...]
invalid_move = [] # [[0,1], ...]
):
global data
global stop_is_pressed
global type_process
out, source, weights, view_img, save_txt, imgsz = \
output, source, weights, view_img, save_txt, img_size
webcam = source == '0' or source.startswith(
'rtsp') or source.startswith('http') or source.endswith('.txt')
# initialize deepsort
cfg = get_config()
cfg.merge_from_file('./deep_sort/configs/deep_sort.yaml')
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
# Initialize
device = torch_utils.select_device(device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
google_utils.attempt_download(weights)
model = torch.load(weights, map_location=device)[
'model'].float().eval() # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
# view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
# dataset = LoadWebcam(source, img_size=imgsz)
else:
save_img = True
# view_img = True
dataset = LoadImages(source, img_size=imgsz)
# dataset = LoadStreams(source, img_size=imgsz)
# save_img = True
# view_img = True
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
# run once
_ = model(img.half() if half else img) if device.type != 'cpu' else None
# print(model)
k = 0
limit = 60
id_limit = 50
trap_xy = 0
trap_frame = False
output_all_frames = {}
counting_id = []
invalid_direction_id = []
invalid_turn_id = []
for path, img, im0s, vid_cap in dataset:
# print(stop_is_pressed)
while(stop_is_pressed):
# cv2.imwrite('./static/images/last_im0.jpg', im0)
yield (b'--frame\r\n'
b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')
# break
# else:
k += 1
trap_xy += 1
if k >= fps_count:
k = 0
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# print(model(img, augment=augment)[0])
# Inference
t1 = torch_utils.time_synchronized()
pred = model(img, augment=augment)[0]
# print('pred b4 nms', pred)
# Apply NMS
pred = non_max_suppression(
pred, conf_thres, iou_thres, classes=classes, agnostic=agnostic_nms)
t2 = torch_utils.time_synchronized()
# print('pred', pred)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
txt_path = str(Path(out) / Path(p).stem) + ('_%g' %
dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
# normalization gain whwh
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(
img.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
bbox_xywh = []
confs = []
save_from_det = {}
for x in range(0, 5):
save_from_det[x] = []
# Write results
for *xyxy, conf, cls in det:
for x in range(0, 5):
if int(cls) == x:
# data[x] += 1
save_from_det[x].append(
[int(xyxy[0].item()), int(xyxy[1].item()),
int(xyxy[2].item()), int(xyxy[3].item())])
break
img_h, img_w, _ = im0.shape
x_c, y_c, bbox_w, bbox_h = bbox_rel(img_w, img_h, *xyxy)
obj = [x_c, y_c, bbox_w, bbox_h]
bbox_xywh.append(obj)
confs.append([conf.item()])
# if save_txt: # Write to file
# xywh = (xyxy2xywh(torch.tensor(xyxy).view(
# 1, 4)) / gn).view(-1).tolist() # normalized xywh
# with open(txt_path + '.txt', 'a') as f:
# f.write(('%g ' * 5 + '\n') %
# (cls, *xywh)) # label format
# if save_img or view_img: # Add bbox to image
# label = '%s %.2f' % (names[int(cls)], conf)
# plot_one_box(xyxy, im0, label=label,
# color=colors[int(cls)], line_thickness=2)
# print('Save from det : ', save_from_det)
xywhs = torch.Tensor(bbox_xywh)
confss = torch.Tensor(confs)
for p, q in line_coordinat:
cv2.line(im0, p, q, (100, 255, 100), 2)
for x in range(len(polygon)):
poly = polygon[x][0]
pts = np.array(poly, np.int32)
pts = pts.reshape((-1, 1, 2))
cv2.polylines(im0, [pts], True, (255, 0, 0), 1)
# Pass detections to deepsort
outputs = deepsort.update(xywhs, confss, im0)
# print('Output Deep Sort: ', outputs)
# draw boxes for visualization
if len(outputs) > 0:
bbox_xyxy = outputs[:, :4]
identities = outputs[:, -1]
draw_boxes(im0, bbox_xyxy, identities)
# Save all results to dictionary
for i, box in enumerate(bbox_xyxy):
x1, y1, x2, y2 = [int(i) for i in box]
# print('x1 y1 x2 y2 : ', x1,y1, x2, y2)
# print('i : ', i)
# print('int(identities[i]) : ', int(identities[i]))
ds_class = float('inf')
smallest = float('inf')
for x in save_from_det:
for sx1, sy1, sx2, sy2 in save_from_det[x]:
diff = sum(abs(np.array([x1, y1, x2, y2])
-np.array([sx1, sy1, sx2, sy2])))
if diff < smallest:
smallest = diff
ds_class = x
if int(identities[i]) in output_all_frames.keys():
# check crossed line
if type_process[0]:
(w1, h1) = (x2-x1, y2-y1)
prev_xyxy = output_all_frames[int(identities[i])][0][-1]
# print('prev xyxy', prev_xyxy)
(xp, yp) = (int(prev_xyxy[0]), int(prev_xyxy[1]))
(wp, hp) = (int(prev_xyxy[2]-xp), int(prev_xyxy[3]-yp))
# p1 = (int(x1 + (w1-x1)/2), int(y1 + (h1-y1)/2))
# q1 = (int(xp + (wp-xp)/2), int(yp + (hp-yp)/2))
p1 = (int(x1 + (w1)/2), int(y1 + (h1)/2))
q1 = (int(xp + (wp)/2), int(yp + (hp)/2))
# print('p1 q1 : ', p1, q1)
cv2.line(im0, p1, q1, (10, 255, 10), 3)
pt1 = IPoint(p1[0], p1[1])
qt1 = IPoint(q1[0], q1[1])
for p2, q2 in line_coordinat:
p2 = IPoint(p2[0], p2[1])
q2 = IPoint(q2[0], q2[1])
if doIntersect(pt1, qt1, p2, q2):
if int(identities[i]) not in counting_id:
counting_id.append(int(identities[i]))
data[most_frequent(output_all_frames[int(identities[i])][1])] += 1
# check direction
# if type_process[1] and len(output_all_frames[int(identities[i])][0]) >= limit:
# #change xyxy to the oldest
# prev_xyxy = output_all_frames[int(identities[i])][0][0]
# (xp, yp) = (int(prev_xyxy[0]), int(prev_xyxy[1]))
# (wp, hp) = (int(prev_xyxy[2]-xp), int(prev_xyxy[3]-yp))
# q1 = (int(xp + (wp)/2), int(yp + (hp)/2))
# minus_x = q1[0] - p1[0]
# minus_y = q1[1] - p1[1]
# # minus_y1 = prev_xyxy[1] - y1
# # minus_y2 = prev_xyxy[3] - y2
# # minus_x1 = prev_xyxy[0] - x1
# # minus_x2 = prev_xyxy[2] - x2
# # 0=up, 1=right, 2=down, 3=left
# # if minus_y1 > 0 and minus_y2 > 0:
# if minus_y > 0:
# output_all_frames[int(identities[i])][4].append(0)
# label = '^'
# t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1.2, 1)[0]
# cv2.putText(im0, label, (x1, y1 - int(t_size[1]/2)), cv2.FONT_HERSHEY_PLAIN, 1.2, [255, 255, 255], 1)
# # if minus_y1 < 0 and minus_y2 < 0:
# if minus_y < 0:
# output_all_frames[int(identities[i])][4].append(2)
# label = 'v'
# t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1.2, 1)[0]
# cv2.putText(im0, label, (x1, y1 - int(t_size[1]/2)), cv2.FONT_HERSHEY_PLAIN, 1.2, [255, 255, 255], 1)
# # if minus_x1 > 0 and minus_x2 > 0:
# if minus_x > 0:
# output_all_frames[int(identities[i])][4].append(3)
# label = '<'
# t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1.2, 1)[0]
# cv2.putText(im0, label, (x1, y1 - int(t_size[1]/2)), cv2.FONT_HERSHEY_PLAIN, 1.2, [255, 255, 255], 1)
# # if minus_x1 < 0 and minus_x2 < 0:
# if minus_x < 0:
# output_all_frames[int(identities[i])][4].append(1)
# label = '>'
# t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1.2, 1)[0]
# cv2.putText(im0, label, (x1, y1 - int(t_size[1]/2)), cv2.FONT_HERSHEY_PLAIN, 1.2, [255, 255, 255], 1)
# check region
l_check = 3
if type_process[2]:
for n in range(len(polygon)):
path = MPath.Path(polygon[n][0])
# inside2 = path.contains_points([[i[0], i[1]]])
if path.contains_point((x1+int(w1/2), y1+int(h1/2))):
output_all_frames[int(identities[i])][2].append(n+1)
output_all_frames[int(identities[i])][3].append(polygon[n][1])
#check direction
# if type_process[1] and len(output_all_frames[int(identities[i])][0]) >= limit:
if type_process[1] and len(output_all_frames[int(identities[i])][0]) >= 1/3*limit\
and (trap_xy >= 10 or trap_frame):
trap_xy = 0
trap_frame = True
# change xyxy to the oldest
prev_xyxy = output_all_frames[int(identities[i])][0][0]
(xp, yp) = (int(prev_xyxy[0]), int(prev_xyxy[1]))
(wp, hp) = (int(prev_xyxy[2]-xp), int(prev_xyxy[3]-yp))
q1 = (int(xp + (wp)/2), int(yp + (hp)/2))
minus_x = q1[0] - p1[0]
minus_y = q1[1] - p1[1]
limit_dir = 1/3
if minus_y > 0 and abs(minus_y) > abs(limit_dir*minus_x):
output_all_frames[int(identities[i])][4].append(0)
label = '^'
t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1.2, 1)[0]
cv2.putText(im0, label, (x1, y1 - int(t_size[1]/2)), cv2.FONT_HERSHEY_PLAIN, 1.2, [255, 255, 255], 1)
if minus_y < 0 and abs(minus_y) > abs(limit_dir*minus_x):
output_all_frames[int(identities[i])][4].append(2)
label = 'v'
t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1.2, 1)[0]
cv2.putText(im0, label, (x1, y1 - int(t_size[1]/2)), cv2.FONT_HERSHEY_PLAIN, 1.2, [255, 255, 255], 1)
if minus_x > 0 and abs(minus_x) > abs(limit_dir*minus_y):
output_all_frames[int(identities[i])][4].append(3)
label = '<'
t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1.2, 1)[0]
cv2.putText(im0, label, (x1, y1 - int(t_size[1]/2)), cv2.FONT_HERSHEY_PLAIN, 1.2, [255, 255, 255], 1)
if minus_x < 0 and abs(minus_x) > abs(limit_dir*minus_y):
output_all_frames[int(identities[i])][4].append(1)
label = '>'
t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1.2, 1)[0]
cv2.putText(im0, label, (x1, y1 - int(t_size[1]/2)), cv2.FONT_HERSHEY_PLAIN, 1.2, [255, 255, 255], 1)
break
# check for invalid turn
if len(output_all_frames[int(identities[i])][2]) >= int(1/4*limit) and type_process[2]:
# unique, frequency = np.unique(output_all_frames[int(identities[i])][2],
# return_counts=True)
first = True
region_trace = []
for r in output_all_frames[int(identities[i])][2]:
if first:
reg = r
region_trace.append(r)
first = False
if reg != r:
region_trace.append(r)
reg = r
if len(region_trace) > 1:
# reset centroid capture if move to another region
output_all_frames[int(identities[i])][4] = []
for reg1, reg2 in invalid_move:
for k in range(len(region_trace)):
if k+1 >= len(region_trace):
break
if (region_trace[k], region_trace[k+1]) == (reg1, reg2):
if int(identities[i]) not in invalid_turn_id:
invalid_turn_id.append(int(identities[i]))
data[6] += 1
label = 'X'
t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2, 2)[0]
cv2.putText(im0, label, (x1 + int(t_size[1]/2), y1), cv2.FONT_HERSHEY_PLAIN, 2, [0, 0, 255], 2)
# check for invalid direction
if len(output_all_frames[int(identities[i])][3]) >= l_check\
and len(output_all_frames[int(identities[i])][4]) >= l_check\
and type_process[1]:
# if most_frequent(output_all_frames[int(identities[i])][3]) \
# != most_frequent(output_all_frames[int(identities[i])][4]):
unique, frequency = np.unique(output_all_frames[int(identities[i])][4],
return_counts=True)
# true_direction = most_frequent(output_all_frames[int(identities[i])][3])
true_direction = output_all_frames[int(identities[i])][3][-1]
opp_direction = true_direction + 2
if opp_direction > 3:
opp_direction -= 4
id_opp_in_unique = -1
for x in range(len(unique)):
if opp_direction == unique[x]:
id_opp_in_unique = x
break
# if id_opp_in_unique >= 0 and frequency[id_opp_in_unique] > int(1/4*limit):
if id_opp_in_unique >= 0 and frequency[id_opp_in_unique] > l_check:
print(output_all_frames[int(identities[i])][4], int(identities[i]))
if int(identities[i]) not in invalid_direction_id:
invalid_direction_id.append(int(identities[i]))
data[5] += 1
label = '!'
t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2, 2)[0]
cv2.putText(im0, label, (x1 + int(t_size[1]/2), y1), cv2.FONT_HERSHEY_PLAIN, 2, [0, 0, 255], 2)
# if true_direction not in unique:
# if int(identities[i]) not in invalid_direction_id:
# invalid_direction_id.append(int(identities[i]))
# data[5] += 1
# label = '!'
# t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2, 2)[0]
# cv2.putText(im0, label, (x1 + int(t_size[1]/2), y1), cv2.FONT_HERSHEY_PLAIN, 2, [0, 0, 255], 2)
# else:
# for x in range(len(unique)):
# if true_direction == unique[x]:
# id_true_in_unique = x
# break
# if frequency[id_true_in_unique] < int(1/3*limit):
# if int(identities[i]) not in invalid_direction_id:
# invalid_direction_id.append(int(identities[i]))
# data[5] += 1
# label = '!'
# t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2, 2)[0]
# cv2.putText(im0, label, (x1 + int(t_size[1]/2), y1), cv2.FONT_HERSHEY_PLAIN, 2, [0, 0, 255], 2)
else:
# oaf[ID] = [[in frame coordinat], [class_type], [region]
# [true_direction], [pred_direction]]
output_all_frames[int(identities[i])] = [[], [], [], [], []]
output_all_frames[int(identities[i])][0].append((x1, y1, x2, y2))
if len(output_all_frames[int(identities[i])][0]) > limit:
output_all_frames[int(identities[i])][0] = output_all_frames[int(identities[i])][0][-limit:]
output_all_frames[int(identities[i])][1].append(ds_class)
if len(output_all_frames[int(identities[i])][1]) > limit:
output_all_frames[int(identities[i])][1] = output_all_frames[int(identities[i])][1][-limit:]
if len(output_all_frames[int(identities[i])][2]) > limit:
output_all_frames[int(identities[i])][2] = output_all_frames[int(identities[i])][2][-limit:]
if len(output_all_frames[int(identities[i])][3]) > limit:
output_all_frames[int(identities[i])][3] = output_all_frames[int(identities[i])][3][-limit:]
if len(output_all_frames[int(identities[i])][4]) > limit:
output_all_frames[int(identities[i])][4] = output_all_frames[int(identities[i])][4][-limit:]
# delete output_all_frames oldest if more than n number of id
if len(output_all_frames) > id_limit:
unused = list(set(output_all_frames.keys())
-set(sorted(output_all_frames.keys())[-id_limit:]))
for x in unused:
del output_all_frames[x]
# if len(counting_id) > id_limit:
# counting_id = counting_id[-id_limit:]
if len(counting_id) > 3:
counting_id = counting_id[-3:]
# if len(invalid_direction_id) > id_limit:
# invalid_direction_id = invalid_direction_id[-id_limit:]
if len(invalid_direction_id) > 3:
invalid_direction_id = invalid_direction_id[-3:]
# if len(invalid_turn_id) > id_limit:
# invalid_turn_id = invalid_turn_id[-id_limit:]
if len(invalid_turn_id) > 3:
invalid_turn_id = invalid_turn_id[-3:]
# print('All Frame : ', output_all_frames)
# Write MOT compliant results to file
if save_txt and len(outputs) != 0:
for j, output in enumerate(outputs):
bbox_left = output[0]
bbox_top = output[1]
bbox_w = output[2]
bbox_h = output[3]
identity = output[-1]
with open(txt_path, 'a') as f:
f.write(('%g ' * 10 + '\n') % (frame_idx, identity, bbox_left,
bbox_top, bbox_w, bbox_h, -1, -1, -1, -1)) # label format
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
print('saving img!')
cv2.imwrite(save_path, im0)
else:
print('saving video!')
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
vid_writer.write(im0)
# invalid direction count reset
trap_frame = False
# cv2.imshow(p, im0)
# if cv2.waitKey(1) == ord('q'): # q to quit
# raise StopIteration
ret, jpeg = cv2.imencode('.jpg', im0)
frame = jpeg.tobytes()
yield (b'--frame\r\n'
b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
def detect(opt, save_img=False):
out, source, weights, view_img, save_txt, imgsz = opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
global counter
global features
# initialize deepsort
cfg = get_config()
cfg.merge_from_file(opt.config_deepsort)
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST,
min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
# Initialize
device = select_device(opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA. Make faster cmputation with lower precision.
#Write headers into csv file
with open(str(Path(args.output)) + '/results.csv', 'a') as f:
f.write("Time,People Count Changed,TotalCount,ActivePerson,\n")
#Initialize the scheduler for every 2 secs
scheduler = BackgroundScheduler()
scheduler.start()
scheduler.add_job(func=write_csv,
args=[opt.output],
trigger=IntervalTrigger(seconds=2))
# Load model
model = torch.load(weights,
map_location=device)['model'].float() # load to FP32
model.to(device).eval()
if half:
model.half() # to FP16
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
view_img = True
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
save_path = str(Path(out))
txt_path = str(Path(out)) + '/results.csv'
for frame_idx, (path, img, im0s, vid_cap) in enumerate(dataset):
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
s += '%gx%g ' % img.shape[2:] # print string
save_path = str(Path(out) / Path(p).name)
if det is not None and len(det):
# Rescale boxes from img_size to im0(640) size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
bbox_xywh = []
confs = []
# Adapt detections to deep sort input format
for *xyxy, conf, cls in det:
img_h, img_w, _ = im0.shape
x_c, y_c, bbox_w, bbox_h = bbox_rel(img_w, img_h, *xyxy)
obj = [x_c, y_c, bbox_w, bbox_h]
bbox_xywh.append(obj)
confs.append([conf.item()])
xywhs = torch.Tensor(bbox_xywh)
confss = torch.Tensor(confs)
# Pass detections to deepsort
outputs = deepsort.update(xywhs, confss, im0)
# draw boxes for visualization
if len(outputs) > 0:
bbox_xyxy = outputs[:, :4]
identities = outputs[:, -1]
draw_boxes(im0, bbox_xyxy, identities)
features['identities'] = identities
if save_txt and len(outputs) != 0:
for j, output in enumerate(outputs):
bbox_left = output[0]
bbox_top = output[1]
bbox_w = output[2]
bbox_h = output[3]
identity = output[-1]
# with open(txt_path, 'a') as f:
# f.write(f"{datetime.now()},{changes if changes != counter else 0},{counter},{len(identities)},\n") # label format
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1), end='\r')
# Write Counter on img
cv2.putText(im0, "Counter : " + str(counter), (10, 20),
cv2.FONT_HERSHEY_PLAIN, 2, [1, 190, 200], 2)
# Stream results
# if view_img:
# cv2.imshow(p, im0)
# if cv2.waitKey(1) == ord('q'): # q to quit
# raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
print('saving img!')
cv2.imwrite(save_path, im0)
else:
# print('saving video!')
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer. Issues with video writer. Fix later
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*opt.fourcc),
fps, (w, h))
vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
print('Done. (%.3fs)' % (time.time() - t0))
def detect(opt, save_img=False):
out, source, weights, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# initialize deepsort
cfg = get_config()
cfg.merge_from_file(opt.config_deepsort)
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST,
min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
# Initialize
device = select_device(opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
#google_utils.attempt_download(weights)
model = torch.load(weights,
map_location=device)['model'].float() # load to FP32
#model = torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning
# model.fuse()
model.to(device).eval()
if half:
model.half() # to FP16
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
view_img = True
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
save_path = str(Path(out))
txt_path = str(Path(out)) + '/results.txt'
for frame_idx, (path, img, im0s, vid_cap) in enumerate(dataset):
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
s += '%gx%g ' % img.shape[2:] # print string
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
bbox_xywh = []
confs = []
# Adapt detections to deep sort input format
for *xyxy, conf, cls in det:
img_h, img_w, _ = im0.shape
x_c, y_c, bbox_w, bbox_h = bbox_rel(img_w, img_h, *xyxy)
obj = [x_c, y_c, bbox_w, bbox_h]
bbox_xywh.append(obj)
confs.append([conf.item()])
xywhs = torch.Tensor(bbox_xywh)
confss = torch.Tensor(confs)
# Pass detections to deepsort
outputs = deepsort.update(xywhs, confss, im0)
# draw boxes for visualization
if len(outputs) > 0:
bbox_xyxy = outputs[:, :4]
identities = outputs[:, -1]
draw_boxes(im0, bbox_xyxy, identities)
# Write MOT compliant results to file
if save_txt and len(outputs) != 0:
for j, output in enumerate(outputs):
bbox_left = output[0]
bbox_top = output[1]
bbox_w = output[2]
bbox_h = output[3]
identity = output[-1]
with open(txt_path, 'a') as f:
f.write(('%g ' * 10 + '\n') %
(frame_idx, identity, bbox_left, bbox_top,
bbox_w, bbox_h, -1, -1, -1,
-1)) # label format
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*opt.fourcc),
fps, (w, h))
vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
def detect(opt, save_img=False):
out, source, weights, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
array_detected_object = []
# initialize deepsort
cfg = get_config()
cfg.merge_from_file(opt.config_deepsort)
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST,
min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
# Initialize
device = select_device(opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = torch.load(weights,
map_location=device)['model'].float() # load to FP32
model.to(device).eval()
if half:
model.half() # to FP16
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
view_img = True
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
save_path = str(Path(out))
txt_path = str(Path(out)) + '/results.txt'
peopleIn = 0
peopleOut = 0
detectedIds = []
for frame_idx, (path, img, im0s, vid_cap) in enumerate(dataset):
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
s += '%gx%g ' % img.shape[2:] # print string
save_path = str(Path(out) / Path(p).name)
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
bbox_xywh = []
confs = []
# Adapt detections to deep sort input format
for *xyxy, conf, cls in det:
img_h, img_w, _ = im0.shape
x_c, y_c, bbox_w, bbox_h = bbox_rel(img_w, img_h, *xyxy)
obj = [x_c, y_c, bbox_w, bbox_h]
bbox_xywh.append(obj)
confs.append([conf.item()])
xywhs = torch.Tensor(bbox_xywh)
confss = torch.Tensor(confs)
# Pass detections to deepsort
outputs = deepsort.update(xywhs, confss, im0)
# draw boxes for visualization
if len(outputs) > 0:
bbox_xyxy = outputs[:, :4]
identities = outputs[:, -1]
draw_boxes(im0, bbox_xyxy, identities)
# Write MOT compliant results to file
if save_txt and len(outputs) != 0:
for j, output in enumerate(outputs):
bbox_left = output[0]
bbox_top = output[1]
bbox_w = output[2]
bbox_h = output[3]
identity = output[-1]
array_detected_object.append(identity)
array_detected_object = list(
dict.fromkeys(array_detected_object))
xas = 0
yas = 0
if identity >= 0:
xas = bbox_xyxy[0][0]
yas = bbox_xyxy[0][1]
if identity not in detectedIds and int(bbox_top) >= 10:
detectedIds.append(identity)
if int(bbox_top) >= 500 or (int(bbox_left) >= 800
and
int(bbox_top) >= 80):
peopleOut += 1
if int(bbox_top) <= 100:
peopleIn += 1
# with open(txt_path, 'a') as f:
# f.write(('%g ' * 10 + '\n') % (frame_idx, identity, bbox_left,
# bbox_top, bbox_w, bbox_h, -1, -1, -1, -1)) # label format
# f.write(('%g ' * 3 + '\n') % (identity, bbox_left, bbox_top)) # label format
# resultText = str(identity) + '-' + str(bbox_top)
# f.write(resultText + '\n') # label format
# f.write(('%g ' * 4 + '\n') % (-1, frame_idx, -1, -1, str(xas), str(yas))) # label format
# f.write(str(identity))
# f.write(('%g ' * 1 + '\n') % (identity))
# f.write('\n')
# f.write(str(bbox_xyxy))
# f.write("Number people counted: " + str(len(array_detected_object)))
# with open(txt_path, 'r') as fp:
# line = fp.readline()
# cnt = 1
# while line:
# identity = line.split("-")[0]
# infoCheck = line.split("-")[1]
# if identity not in detectedIds:
# detectedIds.append(identity)
# if int(infoCheck) > 680 :
# peopleOut += 1
# else:
# peopleIn +=1
# print("Line {}: {}".format(cnt, line.strip().split("-")[0]))
# line = fp.readline()
# cnt += 1
# print("All people counted: " + str(peopleIn + peopleOut))
# print("Number people in: " + str(peopleIn))
# print("Number people out: " + str(peopleOut))
font = cv2.FONT_HERSHEY_DUPLEX
cv2.putText(im0,
"People in out counted: " + str(peopleIn + peopleOut),
(50, 100), font, 0.8, (0, 255, 255), 2, font)
cv2.putText(im0, "Number people in: " + str(peopleIn), (50, 135),
font, 0.8, (0, 255, 255), 2, font)
cv2.putText(im0, "Number people out: " + str(peopleOut), (50, 170),
font, 0.8, (0, 255, 255), 2, font)
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
txt_result = str(Path(out)) + '/result-counted.txt'
print("All people counted: " + str(peopleIn + peopleOut))
print("Number people in: " + str(peopleIn))
print("Number people out: " + str(peopleOut))
with open(txt_path, 'a') as f:
f.write("All people counted: " +
str(peopleIn + peopleOut))
f.write("Number people in: " + str(peopleIn))
f.write("Number people out: " + str(peopleOut))
raise StopIteration
# Save results (image with detections)
if save_img:
print('saving img!')
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
print('saving video!')
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*opt.fourcc),
fps, (w, h))
vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
def rtsp_to_mongodb():
with open("/home/asyed/airflow/dags/parameters.json") as f:
parms = json.load(f)
agnostic_nms = parms["agnostic_nms"]
augment = parms["augment"]
classes = parms["classes"]
conf_thres = parms["conf_thres"]
config_deepsort = parms["config_deepsort"]
deep_sort_model = parms["deep_sort_model"]
device = parms["device"]
dnn = False
evaluate = parms["evaluate"]
exist_ok = parms["exist_ok"]
fourcc = parms["fourcc"]
half = False
print(device)
imgsz = parms["imgsz"]
iou_thres = parms["iou_thres"]
max_det = parms["max_det"]
name = parms["name"]
# save_vid = parms["save_vid"]
#show_vid = parms["show_vid"]
source = parms["source"]
visualize = parms["visualize"]
yolo_model = parms["yolo_model"]
webcam = parms["webcam"]
save_txt = parms["save_txt"]
homography = np.array(parms["homography"])
url = "mongodb://localhost:27017"
client = MongoClient(url)
db = client.trajectory_database
today_date = date.today().strftime("%m-%d-%y")
new = "file_image_coordinates_" + today_date
collection = db[new]
cfg = get_config()
cfg.merge_from_file(config_deepsort)
deepsort = DeepSort(deep_sort_model,
max_dist=cfg.DEEPSORT.MAX_DIST,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
device = select_device(device)
half &= device.type != 'cpu' # half precision only supported on CUDA
# The MOT16 evaluation runs multiple inference streams in parallel, each one writing to
# its own .txt file. Hence, in that case, the output folder is not restored
# make new output folder
# Load model
device = select_device(device)
model = DetectMultiBackend(yolo_model, device=device, dnn=dnn)
stride, names, pt, jit, _ = model.stride, model.names, model.pt, model.jit, model.onnx
imgsz = check_img_size(imgsz, s=stride) # check image size
# Half
half &= pt and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
if pt:
model.model.half() if half else model.model.float()
# Set Dataloader
vid_path, vid_writer = None, None
# Check if environment supports image displays
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source,
img_size=imgsz,
stride=stride,
auto=pt and not jit)
bs = len(dataset) # batch_size
vid_path, vid_writer = [None] * bs, [None] * bs
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
if pt and device.type != 'cpu':
model(
torch.zeros(1, 3, *imgsz).to(device).type_as(
next(model.model.parameters()))) # warmup
# global framess_im2
dt, seen = [0.0, 0.0, 0.0, 0.0], 0
# arr = None
past = []
for frame_idx, (path, img, im0s, vid_cap, s) in enumerate(dataset):
t1 = time_sync()
img = torch.from_numpy(img).to(device)
# print("raw_frame",img.shape)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
t2 = time_sync()
dt[0] += t2 - t1
pred = model(img, augment=augment, visualize=visualize)
t3 = time_sync()
dt[1] += t3 - t2
pred = non_max_suppression(pred,
conf_thres,
iou_thres,
classes,
agnostic_nms,
max_det=max_det)
dt[2] += time_sync() - t3
# Process detections
# dets_per_img = []
for i, det in enumerate(pred): # detections per image
seen += 1
if webcam: # batch_size >= 1
p, im0, _ = path[i], im0s[i].copy(), dataset.count
s += f'{i}: '
else:
p, im0, _ = path, im0s.copy(), getattr(dataset, 'frame', 0)
annotator = Annotator(im0, line_width=2, pil=not ascii)
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
xywhs = xyxy2xywh(det[:, 0:4])
confs = det[:, 4]
clss = det[:, 5]
# pass detections to deepsort
t4 = time_sync()
outputs = deepsort.update(xywhs.cpu(), confs.cpu(),
clss.cpu(), im0)
t5 = time_sync()
dt[3] += t5 - t4
if len(outputs) > 0:
for j, (output, conf) in enumerate(zip(outputs,
confs)):
bboxes = output[0:4]
id = output[4]
cls = output[5]
c = int(cls) # integer class
label = f'{id} {names[c]} {conf:.2f}'
annotator.box_label(bboxes,
label,
color=colors(c, True))
if save_txt:
# to MOT format
bbox_left = output[0]
bbox_top = output[1]
bbox_w = output[2] - output[0]
bbox_h = output[3] - output[1]
# bbox_left = bbox_left + bbox_h
bbox_top = bbox_top + bbox_h
agent_data = {
'frame': int(frame_idx + 1),
'agent_id': int(id),
"labels": str(names[c]),
"x": int(bbox_left),
"y": int(bbox_top)
}
print("agent", agent_data)
collection.insert_one(agent_data)
#db.object_detection.insert_one(agent_data)
#db.pedestrian_detection_15_june.insert_one(agent_data)
#db.test_21_july.insert_one(agent_data)
LOGGER.info(
f'{s}Done. YOLO:({t3 - t2:.3f}s), DeepSort:({t5 - t4:.3f}s)'
)
else:
deepsort.increment_ages()
LOGGER.info('No detections')
im0 = annotator.result()
def detect(opt, device, save_img=False):
out, source, weights, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
colorOrder = ['red', 'purple', 'blue', 'green', 'yellow', 'orange']
frame_num = 0
framestr = 'Frame {frame}'
fpses = []
frame_catch_pairs = []
ball_person_pairs = {}
for color in colorDict:
ball_person_pairs[color] = 0
# Read Class Name Yaml
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader)
names = data_dict['names']
# initialize deepsort
cfg = get_config()
cfg.merge_from_file(opt.config_deepsort)
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST,
min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
# Initialize
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
stride = int(model.stride.max()) # model stride
imgsz = check_img_size(imgsz, s=stride) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
classify = False
if classify:
modelc = load_classifier(name='resnet101', n=2) # initialize
modelc.load_state_dict(
torch.load('weights/resnet101.pt',
map_location=device)['model']).to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz, stride=stride)
else:
dataset = LoadImages(source, img_size=imgsz, stride=stride)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))]
# Run inference
if device.type != 'cpu':
model(
torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(
next(model.parameters()))) # run once
t0 = time.time()
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
txt_path = str(Path(out) / Path(p).stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
bbox_xywh = []
confs = []
clses = []
# Write results
for *xyxy, conf, cls in det:
img_h, img_w, _ = im0.shape # get image shape
x_c, y_c, bbox_w, bbox_h = bbox_rel(img_w, img_h, *xyxy)
obj = [x_c, y_c, bbox_w, bbox_h]
bbox_xywh.append(obj)
confs.append([conf.item()])
clses.append([cls.item()])
xywhs = torch.Tensor(bbox_xywh)
confss = torch.Tensor(confs)
clses = torch.Tensor(clses)
# Pass detections to deepsort
outputs = []
global groundtruths_path
if not 'disable' in groundtruths_path:
# print('\nenabled', groundtruths_path)
groundtruths = solution.load_labels(
groundtruths_path, img_w, img_h, frame_num)
if (groundtruths.shape[0] == 0):
outputs = deepsort.update(xywhs, confss, clses, im0)
else:
# print(groundtruths)
xywhs = groundtruths[:, 2:]
tensor = torch.tensor((), dtype=torch.int32)
confss = tensor.new_ones((groundtruths.shape[0], 1))
clses = groundtruths[:, 0:1]
outputs = deepsort.update(xywhs, confss, clses, im0)
if frame_num >= 2:
for real_ID in groundtruths[:, 1:].tolist():
for DS_ID in xyxy2xywh(outputs[:, :5]):
if (abs(DS_ID[0] - real_ID[1]) / img_w < 0.005
) and (abs(DS_ID[1] - real_ID[2]) / img_h <
0.005) and (
abs(DS_ID[2] - real_ID[3]) /
img_w < 0.005) and (
abs(DS_ID[3] - real_ID[4]) /
img_w < 0.005):
id_mapping[DS_ID[4]] = int(real_ID[0])
else:
outputs = deepsort.update(xywhs, confss, clses, im0)
# draw boxes for visualization
if len(outputs) > 0:
bbox_xyxy = outputs[:, :4]
identities = outputs[:, 4]
clses = outputs[:, 5]
scores = outputs[:, 6]
#Temp solution to get correct id's
mapped_id_list = []
for ids in identities:
if (ids in id_mapping):
mapped_id_list.append(int(id_mapping[ids]))
else:
mapped_id_list.append(ids)
ball_detect, frame_catch_pairs, ball_person_pairs = solution.detect_catches(
im0, bbox_xyxy, clses, mapped_id_list, frame_num,
colorDict, frame_catch_pairs, ball_person_pairs,
colorOrder, save_img)
t3 = time_synchronized()
draw_boxes(im0, bbox_xyxy, [names[i] for i in clses],
scores, ball_detect, identities)
else:
t3 = time_synchronized()
#Draw frame number
tmp = framestr.format(frame=frame_num)
t_size = cv2.getTextSize(tmp, cv2.FONT_HERSHEY_PLAIN, 2, 2)[0]
cv2.putText(im0, tmp, (0, (t_size[1] + 10)),
cv2.FONT_HERSHEY_PLAIN, 2, [255, 255, 255], 2)
#Inference Time
fps = (1 / (t3 - t1))
fpses.append(fps)
print('FPS=%.2f' % fps)
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*opt.fourcc),
fps, (w, h))
vid_writer.write(im0)
frame_num += 1
#t4 = time_synchronized()
avgFps = (sum(fpses) / len(fpses))
print('Average FPS = %.2f' % avgFps)
#print('Total Runtime = %.2f' % (t4 - t0))
outpath = os.path.basename(source)
outpath = outpath[:-4]
outpath = out + '/' + outpath + '_out.csv'
solution.write_catches(outpath, frame_catch_pairs, colorOrder)
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
def main():
video_name = '1.avi'
cap = cv2.VideoCapture(f'data/videos/{video_name}')
fource = cv2.VideoWriter_fourcc(*'mp4v')
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(f'runs/track/{video_name}.mp4', fource, 30,
(width, height))
""" yolov5 目标检测器 """
yolov5_detector = YOLOv5Detector(weights='weights/yolov5s.pt',
conf_thres=0.6)
""" deepsort 追踪器 """
cfg = get_config()
cfg.merge_from_file("deep_sort/configs/deep_sort.yaml")
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST,
min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
window_name = 'Anti Running Car Tracking'
# 每个目标记录 target_point_count 个点，据此判断目标是否滞留
target_point_count = 50
path_cars = []
while True:
state, frame = cap.read()
if not state:
break
prev_time = time.time()
""" 检测目标 """
image, bbox_container = yolov5_detector(frame)
""" 仅保留车信息"""
bbox_container = cut_bbox_container(bbox_container)
""" 初始化一些变量 """
xywh_bboxs = []
labels = []
confs = []
for bbox in bbox_container:
xywh_bboxs.append(xyxy_to_xywh(bbox['box']))
labels.append(bbox['class'])
confs.append(bbox['confidence'])
""" 检测到目标后才有追踪 """
if labels:
""" detections --> deepsort """
xywhs = torch.Tensor(xywh_bboxs)
confss = torch.Tensor(confs)
outputs = deepsort.update(xywhs, confss, labels, frame)
obj_ids = []
bbox_draw = []
num = 0
if len(outputs) > 0:
for (x1, y1, x2, y2, label, track_id) in outputs:
bbox_draw.append({'class': label, 'box': [x1, y1, x2, y2]})
obj_ids.append(track_id)
""" 记录所有目标的路径 每个目标记录点数为 target_point_count """
while track_id > len(path_cars):
path_cars.append([])
path_cars[track_id - 1].append(
(0.5 * (x1 + x2), 0.5 * (y1 + y2)))
""" 超过的点数从首点删除 """
while len(path_cars[track_id - 1]) > target_point_count:
path_cars[track_id - 1].remove(path_cars[track_id -
1][0])
""" 绘图显示 """
num = draw_image(frame, bbox_draw, obj_ids, path_cars,
target_point_count)
""" 输出一些信息 """
for info in bbox_draw:
print(info)
print(obj_ids)
print('---')
""" fps"""
fps = int(1 / (time.time() - prev_time))
cv2.putText(frame,
f'fps={fps} max_id={len(path_cars)}', (10, 40),
0,
1, [0, 255, 0],
thickness=1,
lineType=cv2.LINE_AA)
""" 滞留"""
cv2.putText(frame,
f'stop={num}', (10, 80),
0,
1, [0, 255, 0],
thickness=1,
lineType=cv2.LINE_AA)
""" 拥堵"""
if num >= 3:
cv2.putText(frame,
f'crowded{num}', (10, 120),
0,
1, [0, 255, 0],
thickness=1,
lineType=cv2.LINE_AA)
else:
cv2.putText(frame,
f'normal', (10, 120),
0,
1, [0, 255, 0],
thickness=1,
lineType=cv2.LINE_AA)
""" 车的总数"""
cv2.putText(frame,
f'total number={len(bbox_container)}', (10, 160),
0,
1, [0, 255, 0],
thickness=1,
lineType=cv2.LINE_AA)
cv2.imshow(window_name, frame)
vid_writer.write(frame)
cv2.waitKey(1)
""" 点 x 退出 """
if cv2.getWindowProperty(window_name, cv2.WND_PROP_AUTOSIZE) < 1:
break
cap.release()
vid_writer.release()
cv2.destroyAllWindows()
def detect(opt, save_img=False):
out, source, weights, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# initialize deepsort
cfg = get_config()
cfg.merge_from_file(opt.config_deepsort)
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST,
min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
# Initialize
device = select_device(opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = torch.load(weights,
map_location=device)['model'].float() # load to FP32
model.to(device).eval()
if half:
model.half() # to FP16
# Second-stage classifier
classify = False
if classify:
modelc = torch_utils.load_classifier(name='resnet101',
n=2) # initialize
modelc.load_state_dict(
torch.load('weights/resnet101.pt',
map_location=device)['model']) # load weights
modelc.to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
flag = 0 # 비디오 저장시 사용할 플래그
view_img = True
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
idx = -1
compare_dict = {}
# create a new figure or activate an exisiting figure
fig = plt.figure()
ax = fig.add_subplot(111, projection='polar') # 1,1,1그리드
for path, img, im0s, vid_cap in dataset:
#plt
# Plot origin (agent's start point) - 원점=보행자
ax.plot(0, 0, color='black', marker='o', markersize=20, alpha=0.3)
# Plot configuration
ax.set_rticks([])
ax.set_rmax(1)
ax.grid(False)
ax.set_theta_zero_location("S") # 0도가 어디에 있는지-S=남쪽
ax.set_theta_direction(-1) # 시계방향 극좌표
img = torch.from_numpy(img).to(device)
# img 프레임 자르기
# '''input 이미지 프레임 자르기'''
img = img[:, 100:260, :]
temp = img
add_img = temp[:, :, :32]
img = torch.cat((img, add_img), dim=2)
# 결과 이미지 프레임 자르기
#결과 프레임 자르기 (bouding box와 object 매칭 시키기 위해!!)
im0s = im0s[200:520, :, :]
temp = im0s
add_im0s = temp[:, :64, :]
print(add_im0s.shape)
im0s = np.concatenate((im0s, add_im0s), axis=1)
img = img.half() \
if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
idx += 1
if (idx % 10 != 0): # 동영상 길이 유지
if len(outputs) > 0:
ori_im = draw_boxes(im0s, bbox_xyxy, identities,
isCloser) # 이전 정보로 bbox 그리기
vid_writer.write(im0s)
continue
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
p, s, im0 = path, '', im0s #우리는 여기로 감 - 파일경로 전까지의 출력문은 datasets.py에서 삭제해야함
save_path = str(Path(out) / Path(p).name)
#print(dataset.frame) #프레임 번호
#s += '%gx%g ' % img.shape[2:] # print string #영상 사이즈 출력 (예:640x320) - 삭제가능
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
#만약 차량이 detect된 경우
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results #개수와 클래스 출력(예: 5 cars) -삭제가능
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
#s += '%g %ss, ' % (n, names[int(c)]) # add to string
bbox_xywh = []
confs = []
# Adapt detections to deep sort input format
for *xyxy, conf, cls in det:
img_h, img_w, _ = im0.shape #결과프레임의 사이즈
x_c, y_c, bbox_w, bbox_h = bbox_rel(img_w, img_h,
*xyxy) #center좌표, w, h
obj = [x_c, y_c, bbox_w, bbox_h]
bbox_xywh.append(obj)
confs.append([conf.item()])
xywhs = torch.Tensor(bbox_xywh)
confss = torch.Tensor(confs)
# Pass detections to deepsort
outputs = []
if len(bbox_xywh) != 0: #뭔가 detect됐다면 deepsort로 보냄
outputs, coors, frame, bbox_size = deepsort.update(
xywhs, confss, im0, compare_dict, dataset.frame)
# draw boxes for visualization
if len(outputs) > 0:
print("!", outputs)
bbox_xyxy = outputs[:, :4]
identities = outputs[:, 4]
isCloser = outputs[:, -1]
print("isCloser:", isCloser)
print(compare_dict)
ori_im = draw_boxes(im0, bbox_xyxy, identities,
isCloser) # bbox 그리기
alert.show_direction(ax, coors, bbox_size,
isCloser) # 방향 display하는 함수 호출
# Print time (inference + NMS)
#print('%sDone. (%.3fs)' % (s, t2 - t1))
plt.show(block=False)
# plot 차트 저장
# idx가 10이면 1부터 9까지 10과 같은 이미지로 저장
'''
file = '/Users/wonyeong/Desktop/ewha/project/plotimgs/img%d.png' % idx
plt.savefig(file)
for j in range(9):
file = '/Users/wonyeong/Desktop/ewha/project/plotimgs/img%d.png' % (j + idx + 1)
plt.savefig(file)
# 차량이 detect된 경우에만 그린다..
'''
plt.pause(0.01)
plt.cla()
# Stream results
cv2.imshow('frame', im0)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
print(im0s.shape)
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if (flag == 0):
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fps = vid_cap.get(cv2.CAP_PROP_FPS)
#w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
# h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
w = 1344
h = 320
flag = 1
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*opt.fourcc),
fps, (w, h))
else:
vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
def detect(opt, save_img=False):
out, source, weights, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# initialize deepsort
cfg = get_config()
cfg.merge_from_file(opt.config_deepsort)
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST,
min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
# Initialize
device = select_device(opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Read Yaml
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader)
names = data_dict['names']
print(names)
# Load model
#google_utils.attempt_download(weights)
model = torch.load(weights,
map_location=device)['model'].float() # load to FP32
#model = torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning
# model.fuse()
model.to(device).eval()
if half:
model.half() # to FP16
# Second-stage classifier
classify = False
if classify:
modelc = torch_utils.load_classifier(name='resnet101',
n=2) # initialize
modelc.load_state_dict(
torch.load('weights/resnet101.pt',
map_location=device)['model']) # load weights
modelc.to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))]
# Run inference
t_fps = time_synchronized()
frame_fps = 0
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
txt_path = str(Path(out) / Path(p).stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
#print(det[:, -1].unique())
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
bbox_xywh = []
confs = []
#clses = []
# Write results
for *xyxy, conf, cls in det:
img_h, img_w, _ = im0.shape # get image shape
x_c, y_c, bbox_w, bbox_h = bbox_rel(img_w, img_h, *xyxy)
obj = [x_c, y_c, bbox_w, bbox_h]
bbox_xywh.append(obj)
confs.append([conf.item()])
#clses.append([cls.item()])
#outputs, clses2 = deepsort.update((torch.Tensor(bbox_xywh)), (torch.Tensor(confs)), (torch.Tensor(clses)), im0)
#outputs = deepsort.update((torch.Tensor(bbox_xywh)), (torch.Tensor(confs)), im0)
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
xywhs = torch.Tensor(bbox_xywh)
confss = torch.Tensor(confs)
# Pass detections to deepsort
outputs = deepsort.update(xywhs, confss, im0)
# draw boxes for visualization
if len(outputs) > 0:
bbox_tlwh = []
bbox_xyxy = outputs[:, :4]
identities = outputs[:, -1]
ori_im = draw_boxes(im0, bbox_xyxy, identities)
# Print time (inference + NMS)
#print('%sDone. (%.3fs)' % (s, t2 - t1))
print('%sDone.' % s)
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*opt.fourcc),
fps, (w, h))
vid_writer.write(im0)
frame_fps += 1
if ((time_synchronized() - t_fps) >= 1):
print('\n')
print('FPS=%.2f' % (frame_fps))
t_fps = time_synchronized()
frame_fps = 0
#print('FPS=%.2f' % (1/(time_synchronized() - t1)))
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
def detect(opt):
memory = {}
counter = 0
out, source, yolo_model, deep_sort_model, show_vid, save_vid, save_txt, imgsz, evaluate, half, project, name, exist_ok= \
opt.output, opt.source, opt.yolo_model, opt.deep_sort_model, opt.show_vid, opt.save_vid, \
opt.save_txt, opt.imgsz, opt.evaluate, opt.half, opt.project, opt.name, opt.exist_ok
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# initialize deepsort
cfg = get_config()
cfg.merge_from_file(opt.config_deepsort)
deepsort = DeepSort(deep_sort_model,
torch.device("cpu"),
max_dist=cfg.DEEPSORT.MAX_DIST,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET)
# Initialize
device = select_device(opt.device)
half &= device.type != 'cpu' # half precision only supported on CUDA
# The MOT16 evaluation runs multiple inference streams in parallel, each one writing to
# its own .txt file. Hence, in that case, the output folder is not restored
if not evaluate:
if os.path.exists(out):
pass
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
# Directories
save_dir = increment_path(Path(project) / name,
exist_ok=exist_ok) # increment run
save_dir.mkdir(parents=True, exist_ok=True) # make dir
# Load model
device = select_device(device)
model = DetectMultiBackend(yolo_model, device=device, dnn=opt.dnn)
stride, names, pt, jit, _ = model.stride, model.names, model.pt, model.jit, model.onnx
imgsz = check_img_size(imgsz, s=stride) # check image size
# Half
half &= pt and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
if pt:
model.model.half() if half else model.model.float()
# Set Dataloader
vid_path, vid_writer = None, None
# Check if environment supports image displays
if show_vid:
show_vid = check_imshow()
# Dataloader
if webcam:
show_vid = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source,
img_size=imgsz,
stride=stride,
auto=pt and not jit)
bs = len(dataset) # batch_size
else:
dataset = LoadImages(source,
img_size=imgsz,
stride=stride,
auto=pt and not jit)
bs = 1 # batch_size
vid_path, vid_writer = [None] * bs, [None] * bs
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
# extract what is in between the last '/' and last '.'
txt_file_name = source.split('/')[-1].split('.')[0]
txt_path = str(Path(save_dir)) + '/' + txt_file_name + '.txt'
if pt and device.type != 'cpu':
model(
torch.zeros(1, 3, *imgsz).to(device).type_as(
next(model.model.parameters()))) # warmup
dt, seen = [0.0, 0.0, 0.0, 0.0], 0
regionid = set()
for frame_idx, (path, img, im0s, vid_cap, s) in enumerate(dataset):
t1 = time_sync()
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
t2 = time_sync()
dt[0] += t2 - t1
# Inference
visualize = increment_path(save_dir / Path(path).stem,
mkdir=True) if opt.visualize else False
pred = model(img, augment=opt.augment, visualize=visualize)
t3 = time_sync()
dt[1] += t3 - t2
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
opt.classes,
opt.agnostic_nms,
max_det=opt.max_det)
dt[2] += time_sync() - t3
# Process detections
for i, det in enumerate(pred): # detections per image
seen += 1
if webcam: # batch_size >= 1
p, im0, _ = path[i], im0s[i].copy(), dataset.count
s += f'{i}: '
else:
p, im0, _ = path, im0s.copy(), getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # im.jpg, vid.mp4, ...
s += '%gx%g ' % img.shape[2:] # print string
annotator = Annotator(im0,
line_width=2,
font='Arial.ttf',
pil=not ascii)
if det is not None and len(det):
tboxes = []
indexIDs = []
previous = memory.copy()
memory = {}
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
xywhs = xyxy2xywh(det[:, 0:4])
confs = det[:, 4]
clss = det[:, 5]
# pass detections to deepsort
t4 = time_sync()
outputs = deepsort.update(xywhs.cpu(), confs.cpu(), clss.cpu(),
im0)
t5 = time_sync()
dt[3] += t5 - t4
# draw boxes for visualization
if len(outputs) > 0:
for j, (output, conf) in enumerate(zip(outputs, confs)):
bboxes = output[0:4]
id = output[4]
cls = output[5]
roi = [(0, 0), (640, 0), (640, 380), (0, 380)]
(x, y) = (int(bboxes[0]), int(bboxes[1]))
(w, h) = (int(bboxes[2]), int(bboxes[3]))
inside = cv2.pointPolygonTest(np.array(roi), (x, h),
False)
if inside > 0:
regionid.add(id)
c = int(cls) # integer class
label = f' {names[c]} {conf:.2f}'
cv2.putText(im0, "count =" + str(len(regionid)),
(20, 50), 0, 1, (100, 200, 0), 2)
annotator.box_label(bboxes,
label,
color=colors(c, True))
if save_txt:
# to MOT format
bbox_left = output[0]
bbox_top = output[1]
bbox_w = output[2] - output[0]
bbox_h = output[3] - output[1]
# Write MOT compliant results to file
with open(txt_path, 'a') as f:
f.write(('%g ' * 10 + '\n') % (
frame_idx + 1,
id,
bbox_left, # MOT format
bbox_top,
bbox_w,
bbox_h,
-1,
-1,
-1,
-1))
LOGGER.info(
f'{s}Done. YOLO:({t3 - t2:.3f}s), DeepSort:({t5 - t4:.3f}s)'
)
LOGGER.info(f'counter = {len(regionid)}')
else:
deepsort.increment_ages()
LOGGER.info('No detections')
# Stream results
im0 = annotator.result()
if show_vid:
cv2.imshow(str(p), im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_vid:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
if vid_cap: # video
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps,
(w, h))
vid_writer.write(im0)
# Print results
t = tuple(x / seen * 1E3 for x in dt) # speeds per image
LOGGER.info(
f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS, %.1fms deep sort update \
per image at shape {(1, 3, *imgsz)}' % t)
if save_txt or save_vid:
print('Results saved to %s' % save_path)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
def detect(opt):
out, source, yolo_model, deep_sort_model, show_vid, save_vid, save_txt, imgsz, evaluate, half, \
project, exist_ok, update, save_crop = \
opt.output, opt.source, opt.yolo_model, opt.deep_sort_model, opt.show_vid, opt.save_vid, \
opt.save_txt, opt.imgsz, opt.evaluate, opt.half, opt.project, opt.exist_ok, opt.update, opt.save_crop
webcam = source == '0' or source.startswith(
'rtsp') or source.startswith('http') or source.endswith('.txt')
# Initialize
device = select_device(opt.device)
half &= device.type != 'cpu' # half precision only supported on CUDA
# The MOT16 evaluation runs multiple inference streams in parallel, each one writing to
# its own .txt file. Hence, in that case, the output folder is not restored
if not evaluate:
if os.path.exists(out):
pass
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
# Directories
if type(yolo_model) is str: # single yolo model
exp_name = yolo_model.split(".")[0]
elif type(yolo_model) is list and len(yolo_model) == 1: # single models after --yolo_model
exp_name = yolo_model[0].split(".")[0]
else: # multiple models after --yolo_model
exp_name = "ensemble"
exp_name = exp_name + "_" + deep_sort_model.split('/')[-1].split('.')[0]
save_dir = increment_path(Path(project) / exp_name, exist_ok=exist_ok) # increment run if project name exists
(save_dir / 'tracks' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# Load model
model = DetectMultiBackend(yolo_model, device=device, dnn=opt.dnn)
stride, names, pt = model.stride, model.names, model.pt
imgsz = check_img_size(imgsz, s=stride) # check image size
# Half
half &= pt and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
if pt:
model.model.half() if half else model.model.float()
# Set Dataloader
vid_path, vid_writer = None, None
# Check if environment supports image displays
if show_vid:
show_vid = check_imshow()
# Dataloader
if webcam:
show_vid = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt)
nr_sources = len(dataset)
else:
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
nr_sources = 1
vid_path, vid_writer, txt_path = [None] * nr_sources, [None] * nr_sources, [None] * nr_sources
# initialize deepsort
cfg = get_config()
cfg.merge_from_file(opt.config_deepsort)
# Create as many trackers as there are video sources
deepsort_list = []
for i in range(nr_sources):
deepsort_list.append(
DeepSort(
deep_sort_model,
device,
max_dist=cfg.DEEPSORT.MAX_DIST,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
)
)
outputs = [None] * nr_sources
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
# Run tracking
model.warmup(imgsz=(1 if pt else nr_sources, 3, *imgsz)) # warmup
dt, seen = [0.0, 0.0, 0.0, 0.0], 0
for frame_idx, (path, im, im0s, vid_cap, s) in enumerate(dataset):
t1 = time_sync()
im = torch.from_numpy(im).to(device)
im = im.half() if half else im.float() # uint8 to fp16/32
im /= 255.0 # 0 - 255 to 0.0 - 1.0
if len(im.shape) == 3:
im = im[None] # expand for batch dim
t2 = time_sync()
dt[0] += t2 - t1
# Inference
visualize = increment_path(save_dir / Path(path[0]).stem, mkdir=True) if opt.visualize else False
pred = model(im, augment=opt.augment, visualize=visualize)
t3 = time_sync()
dt[1] += t3 - t2
# Apply NMS
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, opt.classes, opt.agnostic_nms, max_det=opt.max_det)
dt[2] += time_sync() - t3
# Process detections
for i, det in enumerate(pred): # detections per image
seen += 1
if webcam: # nr_sources >= 1
p, im0, _ = path[i], im0s[i].copy(), dataset.count
p = Path(p) # to Path
s += f'{i}: '
txt_file_name = p.name
save_path = str(save_dir / p.name) # im.jpg, vid.mp4, ...
else:
p, im0, _ = path, im0s.copy(), getattr(dataset, 'frame', 0)
p = Path(p) # to Path
# video file
if source.endswith(VID_FORMATS):
txt_file_name = p.stem
save_path = str(save_dir / p.name) # im.jpg, vid.mp4, ...
# folder with imgs
else:
txt_file_name = p.parent.name # get folder name containing current img
save_path = str(save_dir / p.parent.name) # im.jpg, vid.mp4, ...
txt_path = str(save_dir / 'tracks' / txt_file_name) # im.txt
s += '%gx%g ' % im.shape[2:] # print string
imc = im0.copy() if save_crop else im0 # for save_crop
annotator = Annotator(im0, line_width=2, pil=not ascii)
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
xywhs = xyxy2xywh(det[:, 0:4])
confs = det[:, 4]
clss = det[:, 5]
# pass detections to deepsort
t4 = time_sync()
outputs[i] = deepsort_list[i].update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0)
t5 = time_sync()
dt[3] += t5 - t4
# draw boxes for visualization
if len(outputs[i]) > 0:
for j, (output, conf) in enumerate(zip(outputs[i], confs)):
bboxes = output[0:4]
id = output[4]
cls = output[5]
if save_txt:
# to MOT format
bbox_left = output[0]
bbox_top = output[1]
bbox_w = output[2] - output[0]
bbox_h = output[3] - output[1]
# Write MOT compliant results to file
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * 10 + '\n') % (frame_idx + 1, id, bbox_left, # MOT format
bbox_top, bbox_w, bbox_h, -1, -1, -1, i))
if save_vid or save_crop or show_vid: # Add bbox to image
c = int(cls) # integer class
label = f'{id} {names[c]} {conf:.2f}'
annotator.box_label(bboxes, label, color=colors(c, True))
if save_crop:
txt_file_name = txt_file_name if (isinstance(path, list) and len(path) > 1) else ''
save_one_box(bboxes, imc, file=save_dir / 'crops' / txt_file_name / names[c] / f'{id}' / f'{p.stem}.jpg', BGR=True)
LOGGER.info(f'{s}Done. YOLO:({t3 - t2:.3f}s), DeepSort:({t5 - t4:.3f}s)')
else:
deepsort_list[i].increment_ages()
LOGGER.info('No detections')
# Stream results
im0 = annotator.result()
if show_vid:
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 1 millisecond
# Save results (image with detections)
if save_vid:
if vid_path[i] != save_path: # new video
vid_path[i] = save_path
if isinstance(vid_writer[i], cv2.VideoWriter):
vid_writer[i].release() # release previous video writer
if vid_cap: # video
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]
save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos
vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
vid_writer[i].write(im0)
# Print results
t = tuple(x / seen * 1E3 for x in dt) # speeds per image
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS, %.1fms deep sort update \
per image at shape {(1, 3, *imgsz)}' % t)
if save_txt or save_vid:
s = f"\n{len(list(save_dir.glob('tracks/*.txt')))} tracks saved to {save_dir / 'tracks'}" if save_txt else ''
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
if update:
strip_optimizer(yolo_model) # update model (to fix SourceChangeWarning)
