345 lines
14 KiB
Python
345 lines
14 KiB
Python
|
|
# pylint: disable=too-many-statements
|
|
|
|
import math
|
|
import glob
|
|
import os
|
|
import copy
|
|
from contextlib import contextmanager
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torchvision
|
|
import matplotlib.pyplot as plt
|
|
from matplotlib.patches import Circle, FancyArrow
|
|
from PIL import Image
|
|
|
|
from .network.process import laplace_sampling
|
|
from .utils import open_annotations
|
|
from .visuals.pifpaf_show import KeypointPainter, image_canvas
|
|
from .network import Loco
|
|
from .network.process import factory_for_gt, preprocess_pifpaf
|
|
|
|
|
|
def social_interactions(idx, centers, angles, dds, stds=None, social_distance=False,
|
|
n_samples=100, threshold_prob=0.25, threshold_dist=2, radii=(0.3, 0.5)):
|
|
"""
|
|
return flag of alert if social distancing is violated
|
|
"""
|
|
xx = centers[idx][0]
|
|
zz = centers[idx][1]
|
|
distances = [math.sqrt((xx - centers[i][0]) ** 2 + (zz - centers[i][1]) ** 2) for i, _ in enumerate(centers)]
|
|
sorted_idxs = np.argsort(distances)
|
|
indices = [idx_t for idx_t in sorted_idxs[1:] if distances[idx_t] <= threshold_dist]
|
|
|
|
# Deterministic
|
|
if n_samples < 2:
|
|
for idx_t in indices:
|
|
if check_f_formations(idx, idx_t, centers, angles,
|
|
radii=radii, # Binary value
|
|
social_distance=social_distance):
|
|
return True
|
|
|
|
# Probabilistic
|
|
else:
|
|
# Samples distance
|
|
dds = torch.tensor(dds).view(-1, 1)
|
|
stds = torch.tensor(stds).view(-1, 1)
|
|
# stds_te = get_task_error(dds) # similar results to MonoLoco but lower true positive
|
|
laplace_d = torch.cat((dds, stds), dim=1)
|
|
samples_d = laplace_sampling(laplace_d, n_samples=n_samples)
|
|
|
|
# Iterate over close people
|
|
for idx_t in indices:
|
|
f_forms = []
|
|
for s_d in range(n_samples):
|
|
new_centers = copy.deepcopy(centers)
|
|
for el in (idx, idx_t):
|
|
delta_d = dds[el] - float(samples_d[s_d, el])
|
|
theta = math.atan2(new_centers[el][1], new_centers[el][0])
|
|
delta_x = delta_d * math.cos(theta)
|
|
delta_z = delta_d * math.sin(theta)
|
|
new_centers[el][0] += delta_x
|
|
new_centers[el][1] += delta_z
|
|
f_forms.append(check_f_formations(idx, idx_t, new_centers, angles,
|
|
radii=radii,
|
|
social_distance=social_distance))
|
|
if (sum(f_forms) / n_samples) >= threshold_prob:
|
|
return True
|
|
return False
|
|
|
|
|
|
def check_f_formations(idx, idx_t, centers, angles, radii, social_distance=False):
|
|
"""
|
|
Check F-formations for people close together:
|
|
1) Empty space of 0.4 + meters (no other people or themselves inside)
|
|
2) People looking inward
|
|
"""
|
|
|
|
# Extract centers and angles
|
|
other_centers = np.array([cent for l, cent in enumerate(centers) if l not in (idx, idx_t)])
|
|
theta0 = angles[idx]
|
|
theta1 = angles[idx_t]
|
|
|
|
# Find the center of o-space as average of two candidates (based on their orientation)
|
|
for radius in radii:
|
|
x_0 = np.array([centers[idx][0], centers[idx][1]])
|
|
x_1 = np.array([centers[idx_t][0], centers[idx_t][1]])
|
|
|
|
mu_0 = np.array([centers[idx][0] + radius * math.cos(theta0), centers[idx][1] - radius * math.sin(theta0)])
|
|
mu_1 = np.array([centers[idx_t][0] + radius * math.cos(theta1), centers[idx_t][1] - radius * math.sin(theta1)])
|
|
o_c = (mu_0 + mu_1) / 2
|
|
|
|
# Verify they are looking inwards.
|
|
# The distance between mus and the center should be less wrt the original position and the center
|
|
d_new = np.linalg.norm(mu_0 - mu_1) / 2 if social_distance else np.linalg.norm(mu_0 - mu_1)
|
|
d_0 = np.linalg.norm(x_0 - o_c)
|
|
d_1 = np.linalg.norm(x_1 - o_c)
|
|
|
|
# Verify no intrusion for third parties
|
|
if other_centers.size:
|
|
other_distances = np.linalg.norm(other_centers - o_c.reshape(1, -1), axis=1)
|
|
else:
|
|
other_distances = 100 * np.ones((1, 1)) # Condition verified if no other people
|
|
|
|
# Binary Classification
|
|
if d_new <= min(d_0, d_1) and np.min(other_distances) > radius:
|
|
return True
|
|
return False
|
|
|
|
|
|
def predict(args):
|
|
|
|
cnt = 0
|
|
args.device = torch.device('cpu')
|
|
if torch.cuda.is_available():
|
|
args.device = torch.device('cuda')
|
|
|
|
# Load data and model
|
|
monoloco = Loco(model=args.model, net='monoloco_pp',
|
|
device=args.device, n_dropout=args.n_dropout, p_dropout=args.dropout)
|
|
|
|
images = []
|
|
images += glob.glob(args.glob) # from cli as a string or linux converts
|
|
|
|
# Option 1: Run PifPaf extract poses and run MonoLoco in a single forward pass
|
|
if args.json_dir is None:
|
|
from .network import PifPaf, ImageList
|
|
pifpaf = PifPaf(args)
|
|
data = ImageList(args.images, scale=args.scale)
|
|
data_loader = torch.utils.data.DataLoader(
|
|
data, batch_size=1, shuffle=False,
|
|
pin_memory=args.pin_memory, num_workers=args.loader_workers)
|
|
|
|
for idx, (image_paths, image_tensors, processed_images_cpu) in enumerate(data_loader):
|
|
images = image_tensors.permute(0, 2, 3, 1)
|
|
|
|
processed_images = processed_images_cpu.to(args.device, non_blocking=True)
|
|
fields_batch = pifpaf.fields(processed_images)
|
|
|
|
# unbatch
|
|
for image_path, image, processed_image_cpu, fields in zip(
|
|
image_paths, images, processed_images_cpu, fields_batch):
|
|
|
|
if args.output_directory is None:
|
|
output_path = image_path
|
|
else:
|
|
file_name = os.path.basename(image_path)
|
|
output_path = os.path.join(args.output_directory, file_name)
|
|
im_size = (float(image.size()[1] / args.scale),
|
|
float(image.size()[0] / args.scale))
|
|
|
|
print('image', idx, image_path, output_path)
|
|
|
|
_, _, pifpaf_out = pifpaf.forward(image, processed_image_cpu, fields)
|
|
|
|
kk, dic_gt = factory_for_gt(im_size, name=image_path, path_gt=args.path_gt)
|
|
image_t = image # Resized tensor
|
|
|
|
# Run Monoloco
|
|
boxes, keypoints = preprocess_pifpaf(pifpaf_out, im_size, enlarge_boxes=False)
|
|
dic_out = monoloco.forward(keypoints, kk)
|
|
dic_out = monoloco.post_process(dic_out, boxes, keypoints, kk, dic_gt, reorder=False)
|
|
|
|
# Print
|
|
show_social(args, image_t, output_path, pifpaf_out, dic_out)
|
|
|
|
print('Image {}\n'.format(cnt) + '-' * 120)
|
|
cnt += 1
|
|
|
|
# Option 2: Load json file of poses from PifPaf and run monoloco
|
|
else:
|
|
for idx, im_path in enumerate(images):
|
|
|
|
# Load image
|
|
with open(im_path, 'rb') as f:
|
|
image = Image.open(f).convert('RGB')
|
|
if args.output_directory is None:
|
|
output_path = im_path
|
|
else:
|
|
file_name = os.path.basename(im_path)
|
|
output_path = os.path.join(args.output_directory, file_name)
|
|
|
|
im_size = (float(image.size[0] / args.scale),
|
|
float(image.size[1] / args.scale)) # Width, Height (original)
|
|
kk, dic_gt = factory_for_gt(im_size, name=im_path, path_gt=args.path_gt)
|
|
image_t = torchvision.transforms.functional.to_tensor(image).permute(1, 2, 0)
|
|
|
|
# Load json
|
|
basename, ext = os.path.splitext(os.path.basename(im_path))
|
|
|
|
extension = ext + '.pifpaf.json'
|
|
path_json = os.path.join(args.json_dir, basename + extension)
|
|
annotations = open_annotations(path_json)
|
|
|
|
# Run Monoloco
|
|
boxes, keypoints = preprocess_pifpaf(annotations, im_size, enlarge_boxes=False)
|
|
dic_out = monoloco.forward(keypoints, kk)
|
|
dic_out = monoloco.post_process(dic_out, boxes, keypoints, kk, dic_gt, reorder=False)
|
|
# Print
|
|
show_social(args, image_t, output_path, annotations, dic_out)
|
|
|
|
print('Image {}\n'.format(cnt) + '-' * 120)
|
|
cnt += 1
|
|
|
|
|
|
def show_social(args, image_t, output_path, annotations, dic_out):
|
|
"""Output frontal image with poses or combined with bird eye view"""
|
|
|
|
assert 'front' in args.output_types or 'bird' in args.output_types, "outputs allowed: front and/or bird"
|
|
|
|
angles = dic_out['angles']
|
|
dds = dic_out['dds_pred']
|
|
stds = dic_out['stds_ale']
|
|
xz_centers = [[xx[0], xx[2]] for xx in dic_out['xyz_pred']]
|
|
|
|
if 'front' in args.output_types:
|
|
|
|
# Resize back the tensor image to its original dimensions
|
|
if not 0.99 < args.scale < 1.01:
|
|
size = (round(image_t.shape[0] / args.scale), round(image_t.shape[1] / args.scale)) # height width
|
|
image_t = image_t.permute(2, 0, 1).unsqueeze(0) # batch x channels x height x width
|
|
image_t = F.interpolate(image_t, size=size).squeeze().permute(1, 2, 0)
|
|
|
|
# Prepare color for social distancing
|
|
colors = ['r' if social_interactions(idx, xz_centers, angles, dds,
|
|
stds=stds,
|
|
threshold_prob=args.threshold_prob,
|
|
threshold_dist=args.threshold_dist,
|
|
radii=args.radii)
|
|
else 'deepskyblue'
|
|
for idx, _ in enumerate(dic_out['xyz_pred'])]
|
|
|
|
# Draw keypoints and orientation
|
|
keypoint_sets, scores = get_pifpaf_outputs(annotations)
|
|
uv_centers = dic_out['uv_heads']
|
|
sizes = [abs(dic_out['uv_heads'][idx][1] - uv_s[1]) / 1.5 for idx, uv_s in
|
|
enumerate(dic_out['uv_shoulders'])]
|
|
keypoint_painter = KeypointPainter(show_box=False)
|
|
|
|
with image_canvas(image_t,
|
|
output_path + '.front.png',
|
|
show=args.show,
|
|
fig_width=10,
|
|
dpi_factor=1.0) as ax:
|
|
keypoint_painter.keypoints(ax, keypoint_sets, colors=colors)
|
|
draw_orientation(ax, uv_centers, sizes, angles, colors, mode='front')
|
|
|
|
if 'bird' in args.output_types:
|
|
with bird_canvas(args, output_path) as ax1:
|
|
draw_orientation(ax1, xz_centers, [], angles, colors, mode='bird')
|
|
draw_uncertainty(ax1, xz_centers, stds)
|
|
|
|
|
|
def get_pifpaf_outputs(annotations):
|
|
"""Extract keypoints sets and scores from output dictionary"""
|
|
if not annotations:
|
|
return [], []
|
|
keypoints_sets = np.array([dic['keypoints'] for dic in annotations]).reshape(-1, 17, 3)
|
|
score_weights = np.ones((keypoints_sets.shape[0], 17))
|
|
score_weights[:, 3] = 3.0
|
|
# score_weights[:, 5:] = 0.1
|
|
# score_weights[:, -2:] = 0.0 # ears are not annotated
|
|
score_weights /= np.sum(score_weights[0, :])
|
|
kps_scores = keypoints_sets[:, :, 2]
|
|
ordered_kps_scores = np.sort(kps_scores, axis=1)[:, ::-1]
|
|
scores = np.sum(score_weights * ordered_kps_scores, axis=1)
|
|
return keypoints_sets, scores
|
|
|
|
|
|
@contextmanager
|
|
def bird_canvas(args, output_path):
|
|
fig, ax = plt.subplots(1, 1)
|
|
fig.set_tight_layout(True)
|
|
output_path = output_path + '.bird.png'
|
|
x_max = args.z_max / 1.5
|
|
ax.plot([0, x_max], [0, args.z_max], 'k--')
|
|
ax.plot([0, -x_max], [0, args.z_max], 'k--')
|
|
ax.set_ylim(0, args.z_max + 1)
|
|
yield ax
|
|
fig.savefig(output_path)
|
|
plt.close(fig)
|
|
print('Bird-eye-view image saved')
|
|
|
|
|
|
def draw_orientation(ax, centers, sizes, angles, colors, mode):
|
|
|
|
if mode == 'front':
|
|
length = 5
|
|
fill = False
|
|
alpha = 0.6
|
|
zorder_circle = 0.5
|
|
zorder_arrow = 5
|
|
linewidth = 1.5
|
|
edgecolor = 'k'
|
|
radiuses = [s / 1.2 for s in sizes]
|
|
else:
|
|
length = 1.3
|
|
head_width = 0.3
|
|
linewidth = 2
|
|
radiuses = [0.2] * len(centers)
|
|
# length = 1.6
|
|
# head_width = 0.4
|
|
# linewidth = 2.7
|
|
radiuses = [0.2] * len(centers)
|
|
fill = True
|
|
alpha = 1
|
|
zorder_circle = 2
|
|
zorder_arrow = 1
|
|
|
|
for idx, theta in enumerate(angles):
|
|
color = colors[idx]
|
|
radius = radiuses[idx]
|
|
|
|
if mode == 'front':
|
|
x_arr = centers[idx][0] + (length + radius) * math.cos(theta)
|
|
z_arr = length + centers[idx][1] + (length + radius) * math.sin(theta)
|
|
delta_x = math.cos(theta)
|
|
delta_z = math.sin(theta)
|
|
head_width = max(10, radiuses[idx] / 1.5)
|
|
|
|
else:
|
|
edgecolor = color
|
|
x_arr = centers[idx][0]
|
|
z_arr = centers[idx][1]
|
|
delta_x = length * math.cos(theta)
|
|
delta_z = - length * math.sin(theta) # keep into account kitti convention
|
|
|
|
circle = Circle(centers[idx], radius=radius, color=color, fill=fill, alpha=alpha, zorder=zorder_circle)
|
|
arrow = FancyArrow(x_arr, z_arr, delta_x, delta_z, head_width=head_width, edgecolor=edgecolor,
|
|
facecolor=color, linewidth=linewidth, zorder=zorder_arrow)
|
|
ax.add_patch(circle)
|
|
ax.add_patch(arrow)
|
|
|
|
|
|
def draw_uncertainty(ax, centers, stds):
|
|
for idx, std in enumerate(stds):
|
|
std = stds[idx]
|
|
theta = math.atan2(centers[idx][1], centers[idx][0])
|
|
delta_x = std * math.cos(theta)
|
|
delta_z = std * math.sin(theta)
|
|
x = (centers[idx][0] - delta_x, centers[idx][0] + delta_x)
|
|
z = (centers[idx][1] - delta_z, centers[idx][1] + delta_z)
|
|
ax.plot(x, z, color='g', linewidth=2.5)
|