def get_pose(img, res_only=False):
threed_points = np.load(
'./pose_references/reference_3d_68_points_trans.npy')
if not isinstance(img, Image.Image):
img = Image.fromarray(img)
transform = transforms.Compose([transforms.ToTensor()])
DEPTH = 18
MAX_SIZE = 1400
MIN_SIZE = 600
POSE_MEAN = "./models/WIDER_train_pose_mean_v1.npy"
POSE_STDDEV = "./models/WIDER_train_pose_stddev_v1.npy"
MODEL_PATH = "./models/img2pose_v1.pth"
pose_mean = np.load(POSE_MEAN)
pose_stddev = np.load(POSE_STDDEV)
img2pose_model = img2poseModel(
DEPTH,
MIN_SIZE,
MAX_SIZE,
pose_mean=pose_mean,
pose_stddev=pose_stddev,
threed_68_points=threed_points,
)
load_model(img2pose_model.fpn_model,
MODEL_PATH,
cpu_mode=str(img2pose_model.device) == "cpu",
model_only=True)
img2pose_model.evaluate()
(w, h) = img.size
image_intrinsics = np.array([[w + h, 0, w // 2], [0, w + h, h // 2],
[0, 0, 1]])
res = img2pose_model.predict([transform(img)])[0]
if res_only:
return res
all_bboxes = res["boxes"].cpu().numpy().astype('float')
poses = []
bboxes = []
threshold = 0.9
for i in range(len(all_bboxes)):
if res["scores"][i] > threshold:
bbox = all_bboxes[i]
pose_pred = res["dofs"].cpu().numpy()[i].astype('float')
pose_pred = pose_pred.squeeze()
print(pose_pred * 180)
poses.append(pose_pred)
bboxes.append(bbox)
return render_plot(img.copy(), poses, bboxes)
def create_model(self, args):
img2pose_model = img2poseModel(
args.depth,
args.min_size[-1],
args.max_size,
pose_mean=self.pose_mean,
pose_stddev=self.pose_stddev,
threed_68_points=self.threed_68_points,
)
load_model(
img2pose_model.fpn_model,
args.pretrained_path,
cpu_mode=str(img2pose_model.device) == "cpu",
model_only=True,
)
img2pose_model.evaluate()
return img2pose_model
def __init__(self, config):
self.config = config
# start tensorboard summary writer
self.writer = SummaryWriter(config.log_path)
# load training dataset generator
if self.config.random_flip or self.config.random_crop:
self.train_loader = LMDBDataLoaderAugmenter(
self.config, self.config.train_source)
else:
self.train_loader = LMDBDataLoader(self.config,
self.config.train_source)
print(f"Training with {len(self.train_loader.dataset)} images.")
# loads validation dataset generator if a validation dataset is given
if self.config.val_source is not None:
self.val_loader = LMDBDataLoader(self.config,
self.config.val_source, False)
# creates model
self.img2pose_model = img2poseModel(
depth=self.config.depth,
min_size=self.config.min_size,
max_size=self.config.max_size,
device=self.config.device,
pose_mean=self.config.pose_mean,
pose_stddev=self.config.pose_stddev,
distributed=self.config.distributed,
gpu=self.config.gpu,
threed_68_points=np.load(self.config.threed_68_points),
threed_5_points=np.load(self.config.threed_5_points),
)
# optimizer for the backbone and heads
if args.optimizer == "Adam":
self.optimizer = optim.Adam(
self.img2pose_model.fpn_model.parameters(),
lr=self.config.lr,
weight_decay=self.config.weight_decay,
)
elif args.optimizer == "SGD":
self.optimizer = optim.SGD(
self.img2pose_model.fpn_model.parameters(),
lr=self.config.lr,
weight_decay=self.config.weight_decay,
momentum=self.config.momentum,
)
else:
raise Exception(
"No optimizer founded, please select between SGD or Adam.")
# loads a model with optimizer so that it can continue training where it stopped
if self.config.resume_path:
print(f"Resuming training from {self.config.resume_path}")
load_model(
self.img2pose_model.fpn_model,
self.config.resume_path,
model_only=False,
optimizer=self.optimizer,
cpu_mode=str(self.config.device) == "cpu",
)
# loads a pretrained model without loading the optimizer
if self.config.pretrained_path:
print(
f"Loading pretrained weights from {self.config.pretrained_path}"
)
load_model(
self.img2pose_model.fpn_model,
self.config.pretrained_path,
model_only=True,
cpu_mode=str(self.config.device) == "cpu",
)
# saves configuration to file for easier retrival later
print(self.config)
self.save_file(self.config, "config.txt")
# saves optimizer config to file for easier retrival later
print(self.optimizer)
self.save_file(self.optimizer, "optimizer.txt")
self.tensorboard_loss_every = max(len(self.train_loader) // 100, 1)
# self.evaluate_every = max(len(self.train_loader) // 1, 1)
# reduce learning rate when the validation loss stops to decrease
if self.config.lr_plateau:
self.scheduler = ReduceLROnPlateau(
self.optimizer,
mode="min",
factor=0.1,
patience=3,
verbose=True,
threshold=0.001,
cooldown=1,
min_lr=0.00001,
)
# stops training before the defined epochs if validation loss stops to decrease
if self.config.early_stop:
self.early_stop = EarlyStop(mode="min", patience=5)
probability_threshold = 0.9
euler_lpf_coeff = 0.33
bb_coord_lpf_coeff = 0.5
np.set_printoptions(suppress=True)
threed_points = np.load('./pose_references/reference_3d_68_points_trans.npy')
transform = transforms.Compose([transforms.ToTensor()])
pose_mean = np.load(POSE_MEAN)
pose_stddev = np.load(POSE_STDDEV)
img2pose_model = img2poseModel(
DEPTH,
MIN_SIZE,
MAX_SIZE,
pose_mean=pose_mean,
pose_stddev=pose_stddev,
threed_68_points=threed_points,
)
load_model(img2pose_model.fpn_model,
MODEL_PATH,
cpu_mode=str(img2pose_model.device) == "cpu",
model_only=True)
img2pose_model.evaluate()
vcap = cv2.VideoCapture(0, cv2.CAP_V4L2)
vcap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'))
vcap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
vcap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
vcap.set(cv2.CAP_PROP_FPS, 30)