google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""From a directory of images and keypoint protobufs, shows images, with kps.
17

18
Invocation:
19
  $ python -m code/show_keypoints <image_dir> [<mesh_file>].
20
"""
21

22
import glob
23
import os
24
import sys
25

26
import cv2
27
import matplotlib.pyplot as plt
28
import numpy as np
29

30
from keypose import utils
31

32
colors = 255 * plt.cm.get_cmap('rainbow')(np.linspace(0, 1.0, 10))[:, :3]
33

34

35
def draw_circle(image, uvd, color, size=3):
36
  # Filled color circle.
37
  cv2.circle(image, (int(uvd[0]), int(uvd[1])), size, color, -1)
38
  # White outline.
39
  cv2.circle(image, (int(uvd[0]), int(uvd[1])), size + 1, (255, 255, 255))
40

41

42
def show_keypoints(image_dir, mesh_file):
43
  """Display keypoints in a keypose data file."""
44
  print('Looking for images in %s' % image_dir)
45
  filenames = glob.glob(os.path.join(image_dir, '*_L.png'))
46
  if not filenames:
47
    print("Couldn't find any PNG files in %s" % image_dir)
48
    exit(-1)
49
  filenames.sort()
50
  print('Found %d files in %s' % (len(filenames), image_dir))
51

52
  obj = None
53
  if mesh_file:
54
    obj = utils.read_mesh(mesh_file)
55

56
  for fname in filenames:
57
    im_l = utils.read_image(fname)
58
    im_r = utils.read_image(fname.replace('_L.png', '_R.png'))
59
    im_mask = utils.read_image(fname.replace('_L.png', '_mask.png'))
60
    im_border = utils.read_image(fname.replace('_L.png', '_border.png'))
61
    cam, _, _, uvds, _, _, transform = utils.read_contents_pb(
62
        fname.replace('_L.png', '_L.pbtxt')
63
    )
64
    print(fname)
65
    ret = show_kps(im_l, im_r, im_mask, im_border, (cam, uvds, transform), obj)
66
    if ret:
67
      break
68

69

70
def show_kps(im_l, im_r, im_border, im_mask, kps, obj=None, size=3):
71
  """Draw left/right images and keypoints using OpenCV."""
72
  cam, uvds, _ = kps
73

74
  im_l = cv2.cvtColor(im_l, cv2.COLOR_BGR2RGB)
75
  im_r = cv2.cvtColor(im_r, cv2.COLOR_BGR2RGB)
76

77
  uvds = np.array(uvds)
78
  for i, uvd in enumerate(uvds):
79
    draw_circle(im_l, uvd, colors[i * 3], size)
80

81
  if obj:
82
    p_matrix = utils.p_matrix_from_camera(cam)
83
    q_matrix = utils.q_matrix_from_camera(cam)
84
    xyzs = utils.project_np(q_matrix, uvds.T)
85
    obj.project_to_uvd(xyzs, p_matrix)
86
    im_l = obj.draw_points(im_l)
87

88
  cv2.imshow('Image Left', im_l)
89
  cv2.imshow('Border', im_border)
90
  cv2.imshow('Mask', im_mask)
91

92
  key = cv2.waitKey()
93
  if key == ord('q'):
94
    return True
95
  return False
96

97

98
def main():
99
  image_dir = sys.argv[1]
100
  if len(sys.argv) < 3:
101
    mesh_file = None
102
  else:
103
    mesh_file = sys.argv[2]
104

105
  show_keypoints(image_dir, mesh_file)
106

107

108
if __name__ == '__main__':
109
  main()
110