Update keypoints editor.

* Parse descriptor type
* Allow octaves to be hidden.
This commit is contained in:
iabdalkader 2017-01-21 23:04:12 +02:00
parent 5302b9004b
commit 387826a4d4

View File

@ -3,27 +3,33 @@
import time import time
import struct import struct
import pygame import pygame
import sys import sys, os
from math import sin, cos, radians from math import sin, cos, radians
SCALE = 3 SCALE = 3
RED = (255,0,0) RED = (255,0,0)
GREEN = (0,255,0) GREEN = (0,255,0)
WHITE = (255,255,255)
GRAY = (127, 127, 127) GRAY = (127, 127, 127)
KEYPOINTS_SIZE = 24 KEYPOINTS_SIZE = 24
def load_keypoints(path): def load_keypoints(path):
kpts = [] kpts = []
max_octave = 0
with open(path, "rb") as f: with open(path, "rb") as f:
desc = struct.unpack("<I", f.read(4))[0]
size = struct.unpack("<I", f.read(4))[0] size = struct.unpack("<I", f.read(4))[0]
for i in range(0, size): for i in range(0, size):
x, y, score, octave, angle = struct.unpack("<HHHHH", f.read(10)) x, y, score, octave, angle = struct.unpack("<HHHHH", f.read(10))
kpts.append([x, y, score, octave, angle, False, f.read(32)]) kpts.append([x, y, score, octave, angle, False, f.read(32)])
return kpts if (octave > max_octave):
max_octave = octave
print("size: %d max octave: %d" %(len(kpts), max_octave))
return (kpts, max_octave)
if __name__ == '__main__': if __name__ == '__main__':
if len(sys.argv)!= 3: if len(sys.argv)!= 2:
print "usage: kpts_editor.py [image] [descriptor]" print "usage: kpts_editor.py [descriptor]"
sys.exit(1) sys.exit(1)
selection = False selection = False
@ -31,8 +37,13 @@ if __name__ == '__main__':
rect = pygame.Rect(0, 0, 0, 0) rect = pygame.Rect(0, 0, 0, 0)
clock = pygame.time.Clock() clock = pygame.time.Clock()
kpts = load_keypoints(sys.argv[2]) desc_path = sys.argv[1]
img = pygame.image.load(sys.argv[1]) img_path = os.path.splitext(desc_path)[0] + '.pgm'
desc_out = os.path.splitext(desc_path)[0] + '_out.orb'
kpts, max_octave = load_keypoints(desc_path)
curr_octave = max_octave
img = pygame.image.load(img_path)
img = pygame.transform.scale(img, (img.get_width()*SCALE, img.get_height()*SCALE)) img = pygame.transform.scale(img, (img.get_width()*SCALE, img.get_height()*SCALE))
screen = pygame.display.set_mode((img.get_width(), img.get_height())) screen = pygame.display.set_mode((img.get_width(), img.get_height()))
pygame.mouse.set_visible(True) pygame.mouse.set_visible(True)
@ -45,6 +56,13 @@ if __name__ == '__main__':
elif event.type == pygame.KEYDOWN: elif event.type == pygame.KEYDOWN:
if event.key == pygame.K_s: if event.key == pygame.K_s:
save_selected = True save_selected = True
elif event.key == pygame.K_1:
if (curr_octave > 1):
curr_octave = curr_octave - 1
elif event.key == pygame.K_2:
if (curr_octave < max_octave):
curr_octave = curr_octave + 1
elif event.type == pygame.MOUSEBUTTONDOWN: elif event.type == pygame.MOUSEBUTTONDOWN:
selection = True selection = True
rect = pygame.Rect(0, 0, 0, 0) rect = pygame.Rect(0, 0, 0, 0)
@ -68,6 +86,9 @@ if __name__ == '__main__':
for kp in kpts: for kp in kpts:
x, y, score, octave, angle, selected, desc = kp x, y, score, octave, angle, selected, desc = kp
if (octave > curr_octave):
continue
x1 = x*SCALE x1 = x*SCALE
y1 = y*SCALE y1 = y*SCALE
angle = radians(angle) angle = radians(angle)
@ -77,7 +98,7 @@ if __name__ == '__main__':
y2 = y1 + cos(angle) * size y2 = y1 + cos(angle) * size
kp[5] = selected = rect.collidepoint((x1, y1)) kp[5] = selected = rect.collidepoint((x1, y1))
color = GREEN if selected else GRAY color = GREEN if selected else WHITE
pygame.draw.circle(screen, color, (x1, y1), size, 1) pygame.draw.circle(screen, color, (x1, y1), size, 1)
pygame.draw.line(screen, color, (x1, y1), (x2, y2)) pygame.draw.line(screen, color, (x1, y1), (x2, y2))
@ -86,7 +107,8 @@ if __name__ == '__main__':
for kp in kpts: for kp in kpts:
if (kp[5]): kpts_sel.append(kp) if (kp[5]): kpts_sel.append(kp)
with open("desc_out.orb", "wb") as f: with open(desc_out, "wb") as f:
f.write(struct.pack("<I", 1))
f.write(struct.pack("<I", len(kpts_sel))) f.write(struct.pack("<I", len(kpts_sel)))
for kp in kpts_sel: for kp in kpts_sel:
x, y, score, octave, angle, selected, desc = kp x, y, score, octave, angle, selected, desc = kp