def run_inference(model, img, world_to_cam, gt_rot=None, output_dir='vis'):
img_t = transform(img)
img_t = img_t.cuda().unsqueeze(0)
H, W, C = img.shape
render_size = (W, H)
pred = model(img_t).detach().cpu().numpy().squeeze()
trans = trans_gt = np.zeros(3)
rot_euler = np.array([0, 0, pred])
#rot_euler = pred
center_projected_pred, axes_projected_pred = proj_axes_from_trans_rot(
trans_gt, rot_euler, render_size)
vis_pred = draw(img.copy(), center_projected_pred, axes_projected_pred)
cv2.putText(vis_pred, "Pred", (20, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
(255, 255, 255), 1)
if gt_rot is not None:
center_projected_gt, axes_projected_gt = proj_axes_from_trans_rot(
trans_gt, rot_euler_gt, render_size)
vis_gt = draw(img.copy(),
center_projected_gt,
axes_projected_gt,
intensity=100)
cv2.putText(vis_gt, "Ground Truth", (20, 20), cv2.FONT_HERSHEY_SIMPLEX,
0.5, (255, 255, 255), 1)
result = np.hstack((vis_gt, vis_pred))
else:
result = vis_pred
return result
def run_inference(model, img, world_to_cam, gt_rot=None, output_dir='vis'):
img_t = transform(img)
img_t = img_t.cuda().unsqueeze(0)
H, W, C = img.shape
render_size = (W, H)
heatmap, pred = model(img_t)
heatmap = heatmap.detach().cpu().numpy()
pred = pred.detach().cpu().numpy().squeeze()
heatmap = heatmap[0][0]
pred_y, pred_x = np.unravel_index(heatmap.argmax(), heatmap.shape)
heatmap = cv2.normalize(heatmap, None, 0, 255,
cv2.NORM_MINMAX).astype(np.uint8)
heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
heatmap = cv2.addWeighted(img, 0.65, heatmap, 0.35, 0)
heatmap = cv2.arrowedLine(heatmap, (100, 100), (pred_x, pred_y), (0, 0, 0),
1)
cv2.putText(heatmap, "Pred Offset", (pred_x, pred_y - 15),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
trans = trans_gt = np.zeros(3)
rot_euler = np.array([0, 0, pred])
#rot_euler = pred
center_projected_pred, axes_projected_pred = proj_axes_from_trans_rot(
trans_gt, rot_euler, render_size)
vis_pred = draw(img.copy(), center_projected_pred, axes_projected_pred)
cv2.putText(vis_pred, "Pred", (20, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
(255, 255, 255), 1)
if gt_rot is not None:
center_projected_gt, axes_projected_gt = proj_axes_from_trans_rot(
trans_gt, rot_euler_gt, render_size)
vis_gt = draw(img.copy(),
center_projected_gt,
axes_projected_gt,
intensity=100)
cv2.putText(vis_gt, "Ground Truth", (20, 20), cv2.FONT_HERSHEY_SIMPLEX,
0.5, (255, 255, 255), 1)
result = np.hstack((vis_gt, vis_pred))
else:
result = vis_pred
result = np.hstack((result, heatmap))
return result
from src.dataset import KeypointsDataset, transform
from src.prediction import Prediction
from datetime import datetime
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
# model
keypoints = Keypoints(NUM_CLASSES, img_height=IMG_HEIGHT, img_width=IMG_WIDTH)
keypoints.load_state_dict(
torch.load('../character_checkpoints/model_1_18_4.pth'))
# cuda
use_cuda = torch.cuda.is_available()
# use_cuda = False
if use_cuda:
torch.cuda.set_device(2)
keypoints = keypoints.cuda()
prediction = Prediction(keypoints, NUM_CLASSES, IMG_HEIGHT, IMG_WIDTH,
IMG_SMALL_HEIGHT, IMG_SMALL_WIDTH, use_cuda)
transform = transform = transforms.Compose([transforms.ToTensor()])
img = Image.open('../data/test_cropped_humans/0.jpg')
img = np.array(img)
img_t = transform(img)
img_t = img_t.cuda()
result, keypoints = prediction.predict(img_t)
keypoints = keypoints.cpu().numpy()
prediction.plot(img, result, keypoints[0])