def inference(args):
@jit.trace(symbolic=False)
def val_func():
pred_boxes = net(net.inputs)
return pred_boxes
# model path
saveDir = config.model_dir
evalDir = config.eval_dir
misc_utils.ensure_dir(evalDir)
model_file = os.path.join(saveDir,
'epoch_{}.pkl'.format(args.resume_weights))
assert os.path.exists(model_file)
# load model
net = network.Network()
net.eval()
check_point = mge.load(model_file)
net.load_state_dict(check_point['state_dict'])
ori_image, image, im_info = get_data(args.img_path)
net.inputs["image"].set_value(image.astype(np.float32))
net.inputs["im_info"].set_value(im_info)
pred_boxes = val_func().numpy()
num_tag = config.num_classes - 1
target_shape = (pred_boxes.shape[0] // num_tag // top_k, top_k)
pred_tags = (np.arange(num_tag) + 1).reshape(-1, 1)
pred_tags = np.tile(pred_tags, target_shape).reshape(-1, 1)
# nms
if if_set_nms:
from set_nms_utils import set_cpu_nms
n = pred_boxes.shape[0] // top_k
idents = np.tile(np.arange(n)[:, None], (1, top_k)).reshape(-1, 1)
pred_boxes = np.hstack((pred_boxes, idents))
keep = pred_boxes[:, -2] > args.thresh
pred_boxes = pred_boxes[keep]
pred_tags = pred_tags[keep]
keep = set_cpu_nms(pred_boxes, 0.5)
pred_boxes = pred_boxes[keep][:, :-1]
pred_tags = pred_tags[keep]
else:
from set_nms_utils import cpu_nms
keep = pred_boxes[:, -1] > args.thresh
pred_boxes = pred_boxes[keep]
pred_tags = pred_tags[keep]
keep = cpu_nms(pred_boxes, 0.5)
pred_boxes = pred_boxes[keep]
pred_tags = pred_tags[keep]
pred_tags = pred_tags.astype(np.int32).flatten()
pred_tags_name = np.array(config.class_names)[pred_tags]
visual_utils.draw_boxes(ori_image, pred_boxes[:, :-1], pred_boxes[:, -1],
pred_tags_name)
name = args.img_path.split('/')[-1].split('.')[-2]
fpath = '/data/jupyter/{}.png'.format(name)
cv2.imwrite(fpath, ori_image)
def plot_image_detections(decodedArray, dataset_path):
for item in decodedArray:
img_path = os.path.join(dataset_path, item['split'], item['video'],
item['fname'])
image = cv2.imread(img_path, cv2.IMREAD_COLOR)
pred_boxes = item['pred_boxes']
merged_pred_boxes = merge_close_detections(pred_boxes)
pred_tags_name = item['tags']
if merged_pred_boxes is not None:
pred_boxes = np.concatenate((pred_boxes, merged_pred_boxes),
axis=0)
pred_tags_name = pred_tags_name = item['tags'] + [
"merged"
] * merged_pred_boxes.shape[0]
print(item)
if pred_boxes.shape[0] != 0:
image = visual_utils.draw_boxes(image,
pred_boxes[:, :4],
scores=pred_boxes[:, 4],
tags=pred_tags_name,
line_thick=1,
line_color='white')
name = img_path.split('/')[-1].split('.')[-2]
# fpath = '{}/{}.png'.format(out_path, name)
# cv2.imwrite(fpath, image)
cv2.imshow(name, image)
key = cv2.waitKey(
800) #pauses for 3 seconds before fetching next image
if key == 27: #if ESC is pressed, exit loop
cv2.destroyAllWindows()
def incremental_linking(start_frame, video_detections, iou_thresh, jumpgap,
plot):
num_frames = len(video_detections)
live_paths = []
dead_paths = []
start_linking = False
for t in range(start_frame, num_frames):
num_boxes = video_detections[t]['pred_boxes'].shape[0]
##look for close persons
if num_boxes == 0: #no persons detected in frame
if not start_linking:
continue
else:
lp_count = path_count(live_paths)
for lp in range(lp_count):
live_paths[lp]['lastfound'] += 1
elif not start_linking: #there are detections
merge_pred_boxes = merge_close_detections(
video_detections[t]['pred_boxes'], only_merged=True)
start_linking = True if len(merge_pred_boxes) > 0 else False
# if start_linking:
# print('Start tracking from {}'.format(t))
if start_linking:
if num_boxes == 0: #no persons detected in frame
lp_count = path_count(live_paths)
for lp in range(lp_count):
live_paths[lp]['lastfound'] += 1
continue
merge_pred_boxes = merge_close_detections(
video_detections[t]['pred_boxes'],
only_merged=False) if num_boxes >= 2 else video_detections[t][
'pred_boxes'] #Join very close persons
# print("video_detections at {}={}".format(t, video_detections[t]['pred_boxes'].shape[0]))
# print("merge_pred_boxes at {}={}".format(t, len(merge_pred_boxes)))
merge_pred_boxes = np.stack(merge_pred_boxes) if isinstance(
merge_pred_boxes, list) else merge_pred_boxes #listo to numpy
num_boxes = merge_pred_boxes.shape[
0] #update number of bboxes in frame
if len(live_paths) == 0: #first frame
for b in range(num_boxes):
live_paths.append({
'frames_name': [video_detections[t]['fname']],
# 'boxes':[video_detections[t]['pred_boxes'][b,:]], ## Tube/list of bboxes
'boxes': [merge_pred_boxes[b, :]],
'len':
1, ##length of tube
'id':
b, #id tube
'foundAt': [t],
'lastfound':
0 #diff between current frame and last frame in path
})
else:
lp_count = path_count(live_paths)
# print('live-paths:',lp_count)
matrices = []
covered_boxes = np.zeros(num_boxes)
for lp in range(lp_count):
# linking_ious = iou_path_box(live_paths[lp], video_detections[t]['pred_boxes'][:,:4], iou_thresh)
linking_ious = iou_path_box(live_paths[lp],
merge_pred_boxes[:, :4],
iou_thresh)
# best_ios = np.argwhere(linking_ious>iou_thresh)[:,1]
# print("linking_ious: ", linking_ious, linking_ious.shape, best_ios.shape, 'argmax:',)
# print("linking_ious: ", linking_ious, linking_ious.shape, 'max indx:', np.argmax(linking_ious), 'max score:', np.max(linking_ious), np.sum(linking_ious))
matrices.append(linking_ious)
# dead_count = 0
for lp in range(lp_count): #check over live tubes
if live_paths[lp][
'lastfound'] < jumpgap: #verify if path is possibly dead
box_to_lp_score = matrices[lp]
if np.sum(
box_to_lp_score
) > iou_thresh: #there is at least on box that match the tube
maxInd = np.argmax(box_to_lp_score)
max_score = np.max(box_to_lp_score)
live_paths[lp]['frames_name'].append(
video_detections[t]['fname'])
live_paths[lp]['len'] += 1
# live_paths[lp]['boxes'].append(video_detections[t]['pred_boxes'][maxInd,:])
live_paths[lp]['boxes'].append(
merge_pred_boxes[maxInd, :])
live_paths[lp]['foundAt'].append(t)
covered_boxes[maxInd] = 1
else:
live_paths[lp]['lastfound'] += 1
# dead_count += 1
## terminate dead paths
live_paths, dead_paths = terminate_paths(live_paths, jumpgap)
## join dead paths
for dp in range(len(dead_paths)):
dead_paths[dp]['id']
live_paths.append(dead_paths[dp])
lp_count = path_count(live_paths)
##start new paths/tubes
if np.sum(covered_boxes) < num_boxes:
for b in range(num_boxes):
if not covered_boxes.flatten()[b]:
live_paths.append({
'frames_name': [video_detections[t]['fname']],
# 'boxes':[video_detections[t]['pred_boxes'][b,:]], ## Tube/list of bboxes
'boxes': [merge_pred_boxes[b, :]],
'len':
1, ##length of tube
'id':
lp_count, #id tube
'foundAt': [t],
'lastfound':
0 #diff between current frame and last frame in path
})
lp_count += 1
# print('Final paths at {}={}'.format(t,len(live_paths)))
# print('Paths ---live_paths[lp][foundAt]=', [lp['foundAt'] for lp in live_paths])
if plot is not None:
dataset_frames_path = plot['dataset_root']
split = video_detections[t]['split']
video = video_detections[t]['video']
frame = video_detections[t]['fname']
img_path = os.path.join(dataset_frames_path, split, video, frame)
image = cv2.imread(img_path, cv2.IMREAD_COLOR)
pred_boxes = video_detections[t]['pred_boxes'] #real bbox
if pred_boxes.shape[0] != 0 and plot:
image = visual_utils.draw_boxes(
image,
pred_boxes[:, :4],
# scores=pred_boxes[:, 4],
# tags=pred_tags_name,
line_thick=1,
line_color='white')
lp = path_count(live_paths)
box_tubes = []
tube_ids = []
for l in range(lp):
# print('---live_paths[lp][foundAt]:', lp)
foundAt = True if t in live_paths[l]['foundAt'] else False
if foundAt:
bbox = live_paths[l]['boxes'][-1]
box_tubes.append(bbox)
tube_ids.append(live_paths[l]['id'])
print('box_tubes:', len(box_tubes))
print('tube_ids:', len(tube_ids), tube_ids)
if len(box_tubes) > 0:
box_tubes = np.array(box_tubes)
image = visual_utils.draw_boxes(
image,
box_tubes[:, :4],
# scores=pred_boxes[:, 4],
ids=tube_ids,
line_thick=2,
line_color='orange')
cv2.imshow('FRAME' + str(t + 1), image)
key = cv2.waitKey(
plot['wait']) #pauses for 3 seconds before fetching next image
if key == 27: #if ESC is pressed, exit loop
cv2.destroyAllWindows()
# print('==>{} Live paths at {} frame'.format(len(live_paths), t))
live_paths = sorted(live_paths, key=lambda i: i['id'])
return live_paths
def tracking(
decodedArray,
vname,
dataset_frames_path="/Users/davidchoqueluqueroman/Documents/DATASETS_Local/RWF-2000/frames",
video_out_path='/Users/davidchoqueluqueroman/Documents/CODIGOS_SOURCES/AVSS2019/videos_',
plot=True):
mot_tracker = Sort(iou_threshold=0.1)
start_tracking = False
tracked_objects = None
if plot:
size = (224, 224)
result = cv2.VideoWriter('{}/{}.avi'.format(video_out_path, vname),
cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'),
10, size)
for i, item in enumerate(decodedArray): #frame by frame
if plot:
img_path = os.path.join(dataset_frames_path, item['split'],
item['video'], item['fname'])
image = cv2.imread(img_path, cv2.IMREAD_COLOR)
# print("image shape:", image.shape)
pred_boxes = item["pred_boxes"]
pred_tags_name = item['tags']
if not start_tracking:
merge_pred_boxes = merge_close_detections(pred_boxes,
only_merged=True)
start_tracking = True if len(merge_pred_boxes) > 0 else False
# pred_boxes = merge_pred_boxes
# print("pred_boxes: ", pred_boxes.shape)
# if merge_pred_boxes is not None:
# print("merge_pred_boxes: ", merge_pred_boxes.shape)
if pred_boxes.shape[0] != 0 and plot:
image = visual_utils.draw_boxes(
image,
pred_boxes[:, :4],
# scores=pred_boxes[:, 4],
# tags=pred_tags_name,
line_thick=1,
line_color='white')
if start_tracking:
if tracked_objects is not None:
if pred_boxes.shape[0] == 0:
pred_boxes = np.empty((0, 5))
print(
'tracked_objects(no persons in frame):frame {}'.format(
i + 1), tracked_objects.shape)
else:
pred_boxes = merge_close_detections(
pred_boxes,
only_merged=False) #merge close bboxes in every frame
pred_boxes = np.stack(pred_boxes)
tracked_objects = mot_tracker.update(pred_boxes)
print('tracked_objects:frame {}'.format(i + 1),
tracked_objects.shape, ' ids: ', tracked_objects[:,
4])
else:
merge_pred_boxes = np.stack(merge_pred_boxes)
tracked_objects = mot_tracker.update(merge_pred_boxes)
print('--tracked_objects:frame {}'.format(i + 1),
tracked_objects.shape, ' ids: ', tracked_objects[:, 4])
if plot:
image = visual_utils.draw_boxes(
image,
tracked_objects[:, :4],
ids=tracked_objects[:, 4],
# tags=["per-"]*tracked_objects.shape[0],
line_thick=2,
line_color='green')
if plot:
result.write(image) # save video
name = img_path.split('/')[-1].split('.')[-2]
cv2.imshow(name, image)
key = cv2.waitKey(
10) #pauses for 3 seconds before fetching next image
if key == 27: #if ESC is pressed, exit loop
cv2.destroyAllWindows()
if plot:
result.release()
# cv2.destroyAllWindows()
return tracked_objects