diff --git a/df3d/core.py b/df3d/core.py index e9e651e..d76641c 100644 --- a/df3d/core.py +++ b/df3d/core.py @@ -295,7 +295,8 @@ def smooth_points2d(self, cam_id, private_cache=dict()): private_cache[cam_id] = smooth_pose2d(cam.points2d) return private_cache[cam_id] - def plot_2d(self, cam_id, img_id, with_corrections=False, smooth=False, joints=[]): + def plot_2d(self, cam_id, img_id, with_corrections=False, + smooth=False, joints=[], reprojection=False): """Plots the 2d pose estimation results. Parameters: @@ -310,13 +311,17 @@ def plot_2d(self, cam_id, img_id, with_corrections=False, smooth=False, joints=[ """ from pyba.config import df3d_bones, df3d_colors - if with_corrections: - pts2d = self.corrected_points2d(cam_id, img_id) - else: - pts2d = None - return self.camNet[cam_id].plot_2d( - img_id, points2d=pts2d, bones=df3d_bones, colors=df3d_colors - ) + if with_corrections and reprojection: + raise ValueError("'with_corrections' and 'reprojection' " + "cannot both be set to True") + + cam = self.camNet[cam_id] + if reprojection: + return cam.plot_reprojections(img_id, self.camNet.points3d, + bones=df3d_bones, colors=df3d_colors) + pts2d = self.corrected_points2d(cam_id, img_id) if with_corrections else None + return cam.plot_2d(img_id, points2d=pts2d, + bones=df3d_bones, colors=df3d_colors) def get_image(self, cam_id, img_id): """Returns the img_id image from cam_id camera.""" diff --git a/df3d/video.py b/df3d/video.py index 5a2d4a4..a8fa59a 100644 --- a/df3d/video.py +++ b/df3d/video.py @@ -110,17 +110,17 @@ def _make_video(video_path, imgs, fps=default_fps): def _resize(current_shape, new_width): width, height = current_shape - ratio = new_width / width; + ratio = new_width / width return (int(width * ratio), int(height * ratio)) -def _compute_2d_img(plot_2d, img_id, cam_id): +def _compute_2d_img(plot_2d, img_id, cam_id, reprojection=True): """Uses plot_2d to generate an image and resizes it using cv2. Returns: A numpy array containing the resized image. """ - img = plot_2d(cam_id, img_id, smooth=True) + img = plot_2d(cam_id, img_id, smooth=True, reprojection=reprojection) img = cv2.resize(img, (img2d_aspect[0]*img3d_dpi, img2d_aspect[1]*img3d_dpi)) return img