├── libs
    ├── __init__.py
    ├── version.py
    ├── ustr.py
    ├── constants.py
    ├── settings.py
    ├── lib.py
    ├── scene.py
    └── canvas.py
├── README.md
├── resources.qrc
├── utils.py
└── annotator.py


/libs/__init__.py:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | python annotator.py
2 | 


--------------------------------------------------------------------------------
/libs/version.py:
--------------------------------------------------------------------------------
1 | __version_info__ = ('1', '5', '2')
2 | __version__ = '.'.join(__version_info__)
3 | 


--------------------------------------------------------------------------------
/libs/ustr.py:
--------------------------------------------------------------------------------
 1 | import sys
 2 | 
 3 | def ustr(x):
 4 |     '''py2/py3 unicode helper'''
 5 | 
 6 |     if sys.version_info < (3, 0, 0):
 7 |         from PyQt4.QtCore import QString
 8 |         if type(x) == str:
 9 |             return x.decode('utf-8')
10 |         if type(x) == QString:
11 |             return unicode(x)
12 |         return x
13 |     else:
14 |         return x  # py3
15 | 


--------------------------------------------------------------------------------
/libs/constants.py:
--------------------------------------------------------------------------------
 1 | SETTING_FILENAME = 'filename'
 2 | SETTING_RECENT_FILES = 'recentFiles'
 3 | SETTING_WIN_SIZE = 'window/size'
 4 | SETTING_WIN_POSE = 'window/position'
 5 | SETTING_WIN_GEOMETRY = 'window/geometry'
 6 | SETTING_LINE_COLOR = 'line/color'
 7 | SETTING_FILL_COLOR = 'fill/color'
 8 | SETTING_ADVANCE_MODE = 'advanced'
 9 | SETTING_WIN_STATE = 'window/state'
10 | SETTING_SAVE_DIR = 'savedir'
11 | SETTING_LAST_OPEN_DIR = 'lastOpenDir'
12 | SETTING_AUTO_SAVE = 'autosave'
13 | SETTING_SINGLE_CLASS = 'singleclass'


--------------------------------------------------------------------------------
/libs/settings.py:
--------------------------------------------------------------------------------
 1 | import pickle
 2 | import os
 3 | import sys
 4 | 
 5 | class Settings(object):
 6 |     def __init__(self):
 7 |         # Be default, the home will be in the same folder as labelImg
 8 |         home = os.path.expanduser("~")
 9 |         self.data = {}
10 |         self.path = os.path.join(home, '.labelImgSettings.pkl')
11 | 
12 |     def __setitem__(self, key, value):
13 |         self.data[key] = value
14 | 
15 |     def __getitem__(self, key):
16 |         return self.data[key]
17 | 
18 |     def get(self, key, default=None):
19 |         if key in self.data:
20 |             return self.data[key]
21 |         return default
22 | 
23 |     def save(self):
24 |         if self.path:
25 |             with open(self.path, 'wb') as f:
26 |                 pickle.dump(self.data, f, pickle.HIGHEST_PROTOCOL)
27 |                 return True
28 |         return False
29 | 
30 |     def load(self):
31 |         if os.path.exists(self.path):
32 |             with open(self.path, 'rb') as f:
33 |                 self.data = pickle.load(f)
34 |                 return True
35 |         return False
36 | 
37 |     def reset(self):
38 |         if os.path.exists(self.path):
39 |             os.remove(self.path)
40 |             print ('Remove setting pkl file ${0}'.format(self.path))
41 |         self.data = {}
42 |         self.path = None


--------------------------------------------------------------------------------
/resources.qrc:
--------------------------------------------------------------------------------
 1 | <!DOCTYPE RCC><RCC version="1.0">
 2 | <qresource>
 3 | 
 4 | <file alias="help">icons/help.png</file>
 5 | <file alias="app">icons/expert2.png</file>
 6 | <file alias="expert">icons/expert2.png</file>
 7 | <file alias="done">icons/done.png</file>
 8 | <file alias="file">icons/file.png</file>
 9 | <file alias="labels">icons/labels.png</file>
10 | <file alias="new">icons/objects.png</file>
11 | <file alias="close">icons/close.png</file>
12 | <file alias="fit-width">icons/fit-width.png</file>
13 | <file alias="fit-window">icons/fit-window.png</file>
14 | <file alias="undo">icons/undo.png</file>
15 | <file alias="hide">icons/eye.png</file>
16 | <file alias="quit">icons/quit.png</file>
17 | <file alias="copy">icons/copy.png</file>
18 | <file alias="edit">icons/edit.png</file>
19 | <file alias="open">icons/open.png</file>
20 | <file alias="save">icons/save.png</file>
21 | <file alias="save-as">icons/save-as.png</file>
22 | <file alias="color">icons/color.png</file>
23 | <file alias="color_line">icons/color_line.png</file>
24 | <file alias="zoom">icons/zoom.png</file>
25 | <file alias="zoom-in">icons/zoom-in.png</file>
26 | <file alias="zoom-out">icons/zoom-out.png</file>
27 | <file alias="delete">icons/cancel.png</file>
28 | <file alias="next">icons/next.png</file>
29 | <file alias="prev">icons/prev.png</file>
30 | <file alias="resetall">icons/resetall.png</file>
31 | <file alias="verify">icons/verify.png</file>
32 | </qresource>
33 | </RCC>
34 | 


--------------------------------------------------------------------------------
/libs/lib.py:
--------------------------------------------------------------------------------
 1 | from math import sqrt, atan2
 2 | from libs.ustr import ustr
 3 | import hashlib
 4 | try:
 5 |     from PyQt5.QtGui import *
 6 |     from PyQt5.QtCore import *
 7 |     from PyQt5.QtWidgets import *
 8 | except ImportError:
 9 |     from PyQt4.QtGui import *
10 |     from PyQt4.QtCore import *
11 | 
12 | 
13 | def newIcon(icon):
14 |     return QIcon(':/' + icon)
15 | 
16 | 
17 | def newButton(text, icon=None, slot=None):
18 |     b = QPushButton(text)
19 |     if icon is not None:
20 |         b.setIcon(newIcon(icon))
21 |     if slot is not None:
22 |         b.clicked.connect(slot)
23 |     return b
24 | 
25 | 
26 | def newAction(parent, text, slot=None, shortcut=None, icon=None,
27 |               tip=None, checkable=False, enabled=True):
28 |     """Create a new action and assign callbacks, shortcuts, etc."""
29 |     a = QAction(text, parent)
30 |     if icon is not None:
31 |         a.setIcon(newIcon(icon))
32 |     if shortcut is not None:
33 |         if isinstance(shortcut, (list, tuple)):
34 |             a.setShortcuts(shortcut)
35 |         else:
36 |             a.setShortcut(shortcut)
37 |     if tip is not None:
38 |         a.setToolTip(tip)
39 |         a.setStatusTip(tip)
40 |     if slot is not None:
41 |         a.triggered.connect(slot)
42 |     if checkable:
43 |         a.setCheckable(True)
44 |     a.setEnabled(enabled)
45 |     return a
46 | 
47 | 
48 | def addActions(widget, actions):
49 |     for action in actions:
50 |         if action is None:
51 |             widget.addSeparator()
52 |         elif isinstance(action, QMenu):
53 |             widget.addMenu(action)
54 |         else:
55 |             widget.addAction(action)
56 | 
57 | 
58 | def labelValidator():
59 |     return QRegExpValidator(QRegExp(r'^[^ \t].+'), None)
60 | 
61 | 
62 | class struct(object):
63 | 
64 |     def __init__(self, **kwargs):
65 |         self.__dict__.update(kwargs)
66 | 
67 | 
68 | def distance(p):
69 |     return sqrt(p.x() * p.x() + p.y() * p.y())
70 | 
71 | def calcAngle(p):
72 |     return atan2(p.y(), p.x())
73 | 
74 | def fmtShortcut(text):
75 |     mod, key = text.split('+', 1)
76 |     return '<b>%s</b>+<b>%s</b>' % (mod, key)
77 | 
78 | 
79 | def generateColorByText(text):
80 |     s = str(ustr(text))   
81 |     hashCode = int(hashlib.sha256(s.encode('utf-8')).hexdigest(), 16)
82 |     r = int((hashCode / 255) % 255)
83 |     g = int((hashCode / 65025)  % 255)
84 |     b = int((hashCode / 16581375)  % 255)
85 |     return QColor(r, g, b, 100)
86 | 


--------------------------------------------------------------------------------
/utils.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | 
  3 | class ColorPalette:
  4 |     def __init__(self, numColors):
  5 |         np.random.seed(1)
  6 | 
  7 |         self.colorMap = np.array([[255, 0, 0],
  8 |                                   [50, 150, 0],
  9 |                                   [0, 0, 255],
 10 |                                   [80, 128, 255],
 11 |                                   [255, 230, 180],
 12 |                                   [255, 0, 255],
 13 |                                   [0, 255, 255],
 14 |                                   [255, 255, 0],
 15 |                                   [0, 255, 0],
 16 |                                   [200, 255, 255],
 17 |                                   [255, 200, 255],
 18 |                                   [100, 0, 0],
 19 |                                   [0, 100, 0],
 20 |                                   [128, 128, 80],
 21 |                                   [0, 50, 128],
 22 |                                   [0, 100, 100],
 23 |                                   [0, 255, 128],
 24 |                                   [0, 128, 255],
 25 |                                   [255, 0, 128],
 26 |                                   [128, 0, 255],
 27 |                                   [255, 128, 0],
 28 |                                   [128, 255, 0],
 29 |         ], dtype=np.uint8)
 30 |         self.colorMap = np.concatenate([self.colorMap, self.colorMap], axis=0)
 31 | 
 32 |         #self.colorMap = np.maximum(self.colorMap, 1)
 33 | 
 34 |         if numColors > self.colorMap.shape[0]:
 35 |             self.colorMap = np.concatenate([self.colorMap, np.random.randint(255, size = (numColors, 3), dtype=np.uint8)], axis=0)
 36 |             pass
 37 | 
 38 |         return
 39 | 
 40 |     def getColorMap(self):
 41 |         return self.colorMap
 42 | 
 43 |     def getColor(self, index):
 44 |         if index >= colorMap.shape[0]:
 45 |             return np.random.randint(255, size = (3), dtype=np.uint8)
 46 |         else:
 47 |             return self.colorMap[index]
 48 |             pass
 49 | 
 50 | def intersectFaceLine(face, line, return_ratio=False):
 51 |     faceNormal = np.cross(face[1] - face[0], face[2] - face[0])
 52 |     faceArea = 0
 53 |     for c in xrange(1, len(face) - 1):
 54 |         faceArea += np.linalg.norm(np.cross(face[c] - face[0], face[c + 1] - face[c])) / 2
 55 |         pass
 56 |     faceNormal /= np.maximum(faceArea * 2, 1e-4)
 57 |     faceOffset = np.sum(faceNormal * face[0])
 58 |     offset_1 = np.sum(faceNormal * line[0])
 59 |     offset_2 = np.sum(faceNormal * line[1])
 60 |     if offset_2 == offset_1:
 61 |         if return_ratio:
 62 |             return False, 0
 63 |         else:
 64 |             return False
 65 | 
 66 |     alpha = (faceOffset - offset_1) / (offset_2 - offset_1)
 67 |     if alpha <= 0 or alpha >= 1:
 68 |         if return_ratio:
 69 |             return False, alpha
 70 |         else:
 71 |             return False
 72 | 
 73 |     point = line[0] + alpha * (line[1] - line[0])
 74 |     intersectionArea = 0
 75 |     for c in xrange(len(face)):
 76 |         intersectionArea += np.linalg.norm(np.cross(point - face[c], point - face[(c + 1) % len(face)])) / 2
 77 |         continue
 78 |     #print(intersectionArea, faceArea)
 79 |     if intersectionArea <= faceArea + 1e-4:
 80 |         if return_ratio:
 81 |             return True, alpha
 82 |         else:
 83 |             return True
 84 |     else:
 85 |         if return_ratio:
 86 |             return False, alpha
 87 |         else:
 88 |             return False
 89 |     return
 90 | 
 91 | 
 92 | if __name__ == '__main__':
 93 |     line = [np.array([ 2.4764291 ,  4.37349266, -9.5168555 ]), np.array([ 2.4764291 ,  4.37349266, 10.4831445 ])]
 94 |     face = [np.array([2.1361478 , 0.01942726, 0.06335368]), np.array([8.41647591, 2.27955277, 0.06335368]), np.array([6.15570054, 8.74293862, 0.06335368]), np.array([-0.12478519,  6.48326369,  0.06335368])]
 95 |     intersection, ratio = intersectFaceLine(face, line, return_ratio=True)
 96 |     print(face)
 97 |     print(line)
 98 |     print(intersection, ratio)
 99 |     exit(1)
100 | 


--------------------------------------------------------------------------------
/annotator.py:
--------------------------------------------------------------------------------
  1 | import sys
  2 | import os
  3 | from functools import partial
  4 | import numpy as np
  5 | 
  6 | try:
  7 |     from PyQt5.QtGui import *
  8 |     from PyQt5.QtCore import *
  9 |     from PyQt5.QtWidgets import *
 10 | except ImportError:
 11 |     # needed for py3+qt4
 12 |     # Ref:
 13 |     # http://pyqt.sourceforge.net/Docs/PyQt4/incompatible_apis.html
 14 |     # http://stackoverflow.com/questions/21217399/pyqt4-qtcore-qvariant-object-instead-of-a-string
 15 |     if sys.version_info.major >= 3:
 16 |         import sip
 17 |         sip.setapi('QVariant', 2)
 18 |     from PyQt4.QtGui import *
 19 |     from PyQt4.QtCore import *
 20 | 
 21 | import resources
 22 | # Add internal libs
 23 | from libs.constants import *
 24 | from libs.lib import struct, newAction, newIcon, addActions, fmtShortcut, generateColorByText
 25 | from libs.settings import Settings
 26 | from libs.canvas import Canvas
 27 | from libs.ustr import ustr
 28 | from libs.version import __version__
 29 | import glob
 30 | 
 31 | __appname__ = 'annotator'
 32 | 
 33 | class MainWindow(QMainWindow):
 34 |     FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
 35 | 
 36 |     def __init__(self):
 37 |         super(MainWindow, self).__init__()
 38 |         self.setWindowTitle(__appname__)
 39 | 
 40 |         self.settings = Settings()
 41 |         self.settings.load()
 42 |         settings = self.settings
 43 | 
 44 |         self.dataFolder = '../floor_plan_chinese/'
 45 | 
 46 |         self.canvas = Canvas()
 47 | 
 48 |         self.setCentralWidget(self.canvas)
 49 | 
 50 |         action = partial(newAction, self)
 51 | 
 52 |         nextU = action('&NextU', self.moveToNextUnannotated,
 53 |                       'n', 'nextU', u'Move to next unannotated example')
 54 |         next = action('&Next', self.moveToNext,
 55 |                       'Ctrl+n', 'next', u'Move to next example')
 56 | 
 57 | 
 58 |         # Store actions for further handling.
 59 |         self.actions = struct(nextU=nextU, next=next)
 60 | 
 61 | 
 62 |         #self.scenePaths = os.listdir(self.dataFolder)
 63 |         imagePaths = glob.glob('../floor_plan_chinese/*') + glob.glob('../floor_plan_chinese/*/*') + glob.glob('../floor_plan_chinese/*/*/*') + glob.glob('../floor_plan_chinese/*/*/*/*')
 64 |         self.imagePaths = [imagePath for imagePath in imagePaths if '.jpg' in imagePath or '.png' in imagePath or '.jpeg' in imagePath]
 65 |         print(len(self.imagePaths))
 66 |         self.imageIndex = 0
 67 |         
 68 |         self.moveToNextUnannotated()
 69 | 
 70 |         size = settings.get(SETTING_WIN_SIZE, QSize(640, 480))
 71 |         position = settings.get(SETTING_WIN_POSE, QPoint(0, 0))
 72 |         self.resize(size)
 73 |         self.move(position)
 74 | 
 75 |         self.queueEvent(self.loadImage)
 76 | 
 77 | 
 78 |     def paintCanvas(self):
 79 |         #assert not self.image.isNull(), "cannot paint null image"
 80 |         self.canvas.adjustSize()
 81 |         self.canvas.update()
 82 |         return
 83 | 
 84 |     def moveToNextUnannotated(self):
 85 |         self.imageIndex = (self.imageIndex + 1) % len(self.imagePaths)
 86 |         self.loadImage()
 87 |         return
 88 | 
 89 |     def moveToNext(self):
 90 |         self.imageIndex = (self.imageIndex + 1) % len(self.imagePaths)
 91 |         self.loadImage()        
 92 |         return
 93 | 
 94 |     def loadImage(self):
 95 |         imagePath = self.imagePaths[self.imageIndex]
 96 |         self.canvas.loadScene(imagePath)
 97 |         self.paintCanvas()
 98 |         self.setWindowTitle(__appname__ + ' ' + imagePath)
 99 |         self.canvas.setFocus(True)
100 |         return
101 | 
102 | 
103 |     def queueEvent(self, function):
104 |         QTimer.singleShot(0, function)
105 |         return
106 | 
107 | 
108 | def get_main_app(argv=[]):
109 |     """
110 |     Standard boilerplate Qt application code.
111 |     Do everything but app.exec_() -- so that we can test the application in one thread
112 |     """
113 |     app = QApplication(argv)
114 |     app.setApplicationName(__appname__)
115 |     app.setWindowIcon(newIcon("app"))
116 |     # Tzutalin 201705+: Accept extra agruments to change predefined class file
117 |     # Usage : labelImg.py image predefClassFile
118 |     win = MainWindow()
119 |     win.show()
120 |     return app, win
121 | 
122 | 
123 | def main(argv=[]):
124 |     '''construct main app and run it'''
125 |     app, _win = get_main_app(argv)
126 |     return app.exec_()
127 | 
128 | if __name__ == '__main__':
129 |     sys.exit(main(sys.argv))
130 | 


--------------------------------------------------------------------------------
/libs/scene.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | 
  3 | try:
  4 |     from PyQt5.QtGui import *
  5 |     from PyQt5.QtCore import *
  6 | except ImportError:
  7 |     from PyQt4.QtGui import *
  8 |     from PyQt4.QtCore import *
  9 | import copy
 10 | import os
 11 | from utils import *
 12 | import cv2
 13 | import glob
 14 | import json
 15 | 
 16 | COLOR_MAP = [QColor(255, 0, 0), QColor(0, 255, 0), QColor(0, 0, 255), QColor(255, 0, 255)]
 17 | 
 18 | 
 19 | class Scene():
 20 |     def __init__(self, scenePath):
 21 |         self.colorMap = ColorPalette(1000).getColorMap()
 22 |         self.scenePath = scenePath
 23 | 
 24 |         self.imageWidth = 256
 25 |         self.imageHeight = 256
 26 | 
 27 |         image = cv2.imread(scenePath)
 28 |         self.topdownImage = QPixmap(QImage(image.reshape(-1), image.shape[1], image.shape[0], image.shape[1] * 3, QImage.Format_RGB888))
 29 |         os.system('mkdir ' + self.scenePath + '_annotation')
 30 |         self.reset()
 31 |         return
 32 | 
 33 |     def reset(self, mode='load'):
 34 |         self.loading = True
 35 |         #self.topdownImage = None
 36 | 
 37 |         self.corners = []
 38 |         self.connections = []
 39 |         self.prevCornerIndex = -1
 40 | 
 41 |         if os.path.exists(self.scenePath + '_annotation/corners.npy'):
 42 |             self.corners = np.load(self.scenePath + '_annotation/corners.npy').tolist()
 43 |             pass
 44 |         if os.path.exists(self.scenePath + '_annotation/connections.npy'):
 45 |             self.connections = np.load(self.scenePath + '_annotation/connections.npy').tolist()
 46 |             pass
 47 | 
 48 |         #self.fixCornersOnEdges()
 49 |         self.load()
 50 |         return
 51 | 
 52 |     def fixCornersOnEdges(self, epsilon=10):
 53 |         while True:
 54 |             hasChange = False
 55 |             for cornerIndex, corner in enumerate(self.corners):
 56 |                 for connection in self.connections:
 57 |                     if cornerIndex in connection:
 58 |                         continue
 59 |                     corner_1 = np.array(self.corners[connection[0]])
 60 |                     corner_2 = np.array(self.corners[connection[1]])
 61 |                     normal = corner_1 - corner_2
 62 |                     normal /= max(np.linalg.norm(normal), 1e-4)
 63 |                     normal = np.array([normal[1], -normal[0]])
 64 |                     distance = np.dot(corner_1, normal) - np.dot(corner, normal)
 65 |                     if abs(distance) < epsilon * 2 and ((abs(normal[1]) > abs(normal[0]) and corner[0] > min(corner_1[0], corner_2[0]) and corner[0] < max(corner_1[0], corner_2[0])) or (abs(normal[1]) < abs(normal[0]) and corner[1] > min(corner_1[1], corner_2[1]) and corner[1] < max(corner_1[1], corner_2[1]))):
 66 |                         self.connections.append((connection[0], cornerIndex))
 67 |                         self.connections.append((connection[1], cornerIndex))
 68 |                         self.connections.remove(connection)
 69 |                         hasChange = True
 70 |                         break
 71 |                     continue
 72 |                 continue
 73 |             if not hasChange:
 74 |                 break
 75 |             continue
 76 |         return
 77 | 
 78 |     def findRoomCorners(self):
 79 |         #for roomName, roomLabel in self.roomLabelDict.iteritems():
 80 |         self.roomCorners = {}
 81 |         initialCornerSize = 40
 82 |         for cornerIndex, corner in enumerate(self.corners):
 83 |             cornerSize = initialCornerSize
 84 |             corner = np.round(np.array(corner)).astype(np.int32)
 85 |             while True:
 86 |                 roomLabels = self.roomSegmentation[max(corner[1] - cornerSize / 2, 0):min(corner[1] + cornerSize / 2, self.roomSegmentation.shape[0] - 1), max(corner[0] - cornerSize / 2, 0):min(corner[0] + cornerSize / 2, self.roomSegmentation.shape[1] - 1)]
 87 |                 roomLabels = roomLabels[roomLabels > 0]
 88 |                 if len(roomLabels) == 0:
 89 |                     cornerSize += 2
 90 |                     continue
 91 |                 roomLabels = np.unique(roomLabels)
 92 |                 for roomLabel in roomLabels:
 93 |                     if roomLabel not in self.roomCorners:
 94 |                         self.roomCorners[roomLabel] = []
 95 |                         pass
 96 |                     self.roomCorners[roomLabel].append(cornerIndex)
 97 |                     continue
 98 |                 break
 99 |             continue
100 |         return
101 | 
102 | 
103 |     def paint(self, painter, patchOffsets, patchSizes, topdownOffset, topdownScale, offsetX, offsetY):
104 |         if self.loading:
105 |             return
106 |         #sizes = np.array([self.topdown.shape[1], self.topdown.shape[0]])
107 |         #patchOffsets = np.minimum(np.maximum(patchOffsets, 0), sizes - patchSizes)
108 | 
109 |         #topdown = self.topdown[patchOffsets[1]:patchOffsets[1] + patchSizes[1]][patchOffsets[0]:patchOffsets[0] + patchSizes[0]]
110 |         #topdown = np.minimum((topdown - topdownOffset).astype(np.float32) / topdownScale * 255, 255).astype(np.uint8)
111 |         #topdown = np.tile(np.expand_dims(topdown, axis=-1), [1, 1, 3])
112 |         
113 |         painter.drawPixmap(offsetX, offsetY, self.topdownImage)
114 | 
115 |         color = COLOR_MAP[0]
116 |         pen = QPen(color)
117 |         pen.setWidth(3)
118 |         painter.setPen(pen)
119 |         d = 10
120 | 
121 |         corner_path = QPainterPath()
122 |         points = []
123 |         for _, corner in enumerate(self.corners):
124 |             point = QPoint(int(round(corner[0] - patchOffsets[0] + offsetX)), int(round(corner[1] - patchOffsets[1] + offsetY)))
125 |             points.append(point)
126 |             corner_path.addEllipse(point, d / 2.0, d / 2.0)
127 |             continue
128 |         painter.drawPath(corner_path)
129 | 
130 |         connection_path = QPainterPath()
131 |         for connection in self.connections:
132 |             connection_path.moveTo(points[connection[0]])
133 |             connection_path.lineTo(points[connection[1]])
134 |             continue
135 |         painter.drawPath(connection_path)
136 | 
137 | 
138 |         return
139 | 
140 | 
141 |     def addCorner(self, newCorner, axisAligned=True, epsilon=10):
142 |         newCornerIndex = -1
143 |         for cornerIndex, corner in enumerate(self.corners):
144 |             if np.linalg.norm(corner - newCorner) < epsilon:
145 |                 newCornerIndex = cornerIndex
146 |                 break
147 |             continue
148 |         if newCornerIndex == -1:
149 |             newCornerIndex = len(self.corners)
150 |             if self.prevCornerIndex != -1 and axisAligned:
151 |                 delta = newCorner - self.corners[self.prevCornerIndex]
152 |                 if abs(delta[0]) < abs(delta[1]):
153 |                     delta[0] = 0
154 |                 else:
155 |                     delta[1] = 0
156 |                     pass
157 |                 newCorner = self.corners[self.prevCornerIndex] + delta
158 |                 pass
159 | 
160 |             for connection in self.connections:
161 |                 corner_1 = np.array(self.corners[connection[0]])
162 |                 corner_2 = np.array(self.corners[connection[1]])
163 |                 normal = corner_1 - corner_2
164 |                 normal /= max(np.linalg.norm(normal), 1e-4)
165 |                 normal = np.array([normal[1], -normal[0]])
166 |                 distance = np.dot(corner_1, normal) - np.dot(newCorner, normal)
167 |                 #print(abs(normal[1]) < abs(normal[0]), corner[1] > min(corner_1[1], corner_2[1]), corner[1] < max(corner_1[1], corner_2[1]))
168 |                 if abs(distance) < epsilon * 2 and ((abs(normal[1]) > abs(normal[0]) and newCorner[0] > min(corner_1[0], corner_2[0]) and newCorner[0] < max(corner_1[0], corner_2[0])) or (abs(normal[1]) < abs(normal[0]) and newCorner[1] > min(corner_1[1], corner_2[1]) and newCorner[1] < max(corner_1[1], corner_2[1]))):
169 |                     #print(connection, corner_1, corner_2, newCorner, normal)
170 |                     newCorner = newCorner + distance * normal
171 |                     self.connections.append((connection[0], newCornerIndex))
172 |                     self.connections.append((connection[1], newCornerIndex))
173 |                     self.connections.remove(connection)
174 |                     break
175 |                 continue
176 | 
177 |             self.corners.append(newCorner)
178 |             pass
179 |         if self.prevCornerIndex != -1 and self.prevCornerIndex != newCornerIndex:
180 |             self.connections.append((self.prevCornerIndex, newCornerIndex))
181 |             pass
182 |         self.prevCornerIndex = newCornerIndex
183 | 
184 |         return
185 | 
186 |     def moveCorner(self, delta):
187 |         if self.prevCornerIndex == -1:
188 |             return
189 |         self.corners[self.prevCornerIndex] += delta
190 |         pass
191 | 
192 |     def finalize(self):
193 |         self.prevCornerIndex = -1
194 |         return
195 | 
196 |     def save(self):
197 |         #scene_info = {'corners': self.corners, 'cornersOpp': self.cornersOpp, 'faces': self.faces, 'dominantNormals': self.dominantNormals}
198 |         np.save(self.scenePath + '_annotation/corners.npy', np.array(self.corners))
199 |         np.save(self.scenePath + '_annotation/connections.npy', np.array(self.connections))
200 |         return
201 | 
202 |     
203 |     def loadImage(self, imageIndex):
204 |         imagePath = self.imagePaths[imageIndex % len(self.imagePaths)]
205 |         #print(imagePath)
206 |         with open(imagePath.replace('rgb/', 'pose/').replace('rgb.png', 'pose.json')) as f:
207 |             pose = json.load(f)
208 |             pass
209 |         extrinsics = np.array(pose['camera_rt_matrix'])
210 |         intrinsics = np.array(pose['camera_k_matrix'])
211 | 
212 |         image = cv2.imread(imagePath)
213 |         imageSizes = np.array([image.shape[1], image.shape[0]])
214 |         image = cv2.resize(image, (self.imageWidth, self.imageHeight))
215 |         roomName = '_'.join(imagePath.split('/')[-1].split('_')[2:4])
216 |         if roomName not in self.roomLabelDict:
217 |             return None, []
218 |         roomLabel = self.roomLabelDict[roomName]
219 |         if roomLabel not in self.roomCorners:
220 |             return None, []
221 |         cornerIndices = self.roomCorners[roomLabel]
222 |         corners2D = np.array([self.corners[cornerIndex] for cornerIndex in cornerIndices]).astype(np.float32)
223 | 
224 |         #print(corners2D)
225 |         X = corners2D[:, 0] / self.box[4] * (self.box[1] - self.box[0]) + self.box[0]
226 |         Y = -(corners2D[:, 1] / self.box[5] * (self.box[3] - self.box[2]) + self.box[2])
227 |         cornerGT = []
228 |         for cornerType, horizontalHeight in enumerate(self.box[6:8]):
229 |             wallCorners3D = np.stack([X, Y, np.full(X.shape, horizontalHeight), np.ones(X.shape)], axis=-1)
230 |             wallCorners3D = np.matmul(extrinsics, wallCorners3D.transpose()).transpose()
231 |             wallCorners = np.matmul(intrinsics, wallCorners3D.transpose()).transpose()
232 | 
233 |             wallCorners = np.round(wallCorners[:, :2] / wallCorners[:, 2:3] / imageSizes * np.array([self.imageWidth, self.imageHeight])).astype(np.int32)
234 |             cornerMask = np.logical_and(np.logical_and(wallCorners3D[:, 2] > 0, np.logical_and(np.logical_and(np.all(wallCorners >= 0, axis=-1), wallCorners[:, 0] < self.imageWidth), wallCorners[:, 1] < self.imageHeight)), wallCorners3D[:, 2] < 10)
235 | 
236 |             for index, wallCorner in enumerate(wallCorners):
237 |                 if not cornerMask[index]:
238 |                     continue
239 |                 cornerGT.append(wallCorner.tolist() + [cornerType, cornerIndices[index], wallCorners3D[index, 2]])
240 |                 continue
241 |             continue
242 |         cornerGT = np.array(cornerGT)
243 |         if True:
244 |             cornerImage = image.copy()
245 |             for index, wallCorner in enumerate(cornerGT):
246 |                 cv2.circle(cornerImage, (int(round(wallCorner[0])), int(round(wallCorner[1]))), 10, (0, 0, 255), -1)
247 |                 continue
248 |             cv2.imwrite('test/corner.png', cornerImage)
249 |             pass
250 |         return image, cornerGT
251 | 
252 | 
253 |     def loadImages(self):
254 |         self.findRoomCorners()
255 | 
256 |         imagePaths = glob.glob(self.scenePath + '/data/rgb/*.png')
257 |         self.imagePaths = sorted(imagePaths)
258 |         self.trajectory = []
259 |         return
260 |         #centers = []
261 |         for imageIndex, imagePath in enumerate(imagePaths):
262 |             if imageIndex != 80:
263 |                 continue
264 |             self.loadImage(imagePath)
265 |             continue
266 |         #centers = np.array(centers)
267 |         return
268 | 
269 |     def load(self):
270 |         self.loading = False
271 |         return
272 | 
273 |     def removeLast(self):
274 |         if self.prevCornerIndex == len(self.corners) - 1:
275 |             self.corners = self.corners[:-1]
276 |             self.connections = [connection for connection in self.connections if len(self.corners) not in connection]
277 |             self.prevCornerIndex -= 1
278 |         else:
279 |             self.connections = self.connections[:-1]
280 |             pass
281 |         return
282 | 


--------------------------------------------------------------------------------
/libs/canvas.py:
--------------------------------------------------------------------------------
  1 | try:
  2 |     from PyQt5.QtGui import *
  3 |     from PyQt5.QtCore import *
  4 |     from PyQt5.QtWidgets import *
  5 | except ImportError:
  6 |     from PyQt4.QtGui import *
  7 |     from PyQt4.QtCore import *
  8 | 
  9 | #from PyQt4.QtOpenGL import *
 10 | 
 11 | #from libs.corner import Corner
 12 | from libs.lib import distance
 13 | import numpy as np
 14 | import cv2
 15 | from PIL import Image
 16 | import requests
 17 | #import StringIO
 18 | import urllib
 19 | import sys
 20 | import glob
 21 | from scene import Scene
 22 | import os
 23 | 
 24 | sys.path.append('../code/')
 25 | Image.MAX_IMAGE_PIXELS = 1000000000
 26 | 
 27 | CURSOR_DEFAULT = Qt.ArrowCursor
 28 | CURSOR_POINT = Qt.PointingHandCursor
 29 | CURSOR_DRAW = Qt.CrossCursor
 30 | CURSOR_MOVE = Qt.ClosedHandCursor
 31 | CURSOR_GRAB = Qt.OpenHandCursor
 32 | 
 33 | # class Canvas(QGLWidget):
 34 | 
 35 | 
 36 | class Canvas(QWidget):
 37 |     newCorner = pyqtSignal()
 38 |     cornerMoved = pyqtSignal()
 39 |     drawing = pyqtSignal(bool)
 40 | 
 41 |     CREATE, EDIT = list(range(2))
 42 |     image = None
 43 | 
 44 |     def __init__(self, *args, **kwargs):
 45 |         super(Canvas, self).__init__(*args, **kwargs)
 46 |         # Initialise local state.
 47 | 
 48 |         self.prevPoint = QPointF()
 49 |         self._painter = QPainter()
 50 |         # Set widget options.
 51 |         self.setMouseTracking(True)
 52 |         self.setFocusPolicy(Qt.WheelFocus)
 53 | 
 54 |         self.width = 1280
 55 |         self.height = 960
 56 | 
 57 |         self.layout_width = 1000
 58 |         self.layout_height = 1000
 59 | 
 60 |         self.offsetX = 10
 61 |         self.offsetY = 10
 62 | 
 63 |         self.currentLabel = 0
 64 |         self.hiding = False
 65 |         self.resize(self.width, self.height)
 66 |         self.imageIndex = -1
 67 |         self.mode = 'layout'
 68 |         self.ctrlPressed = False
 69 |         self.shiftPressed = False
 70 | 
 71 |         self.patchOffsets = np.zeros(2, dtype=np.int32)
 72 |         self.patchSizes = np.full((2, ), 1000, dtype=np.int32)
 73 |         self.topdownOffset = 0
 74 |         self.topdownScale = 1
 75 |         self.topdownImage = None
 76 |         return
 77 | 
 78 |     # def creating(self):
 79 |     #     return self.mode == self.CREATE
 80 | 
 81 |     # def editing(self):
 82 |     #     return self.mode == self.EDIT
 83 | 
 84 |     # def setMode(self, value=True):
 85 |     #     self.mode = self.CREATE if value else self.EDIT
 86 |     #     if value:  # Create
 87 |     #         self.unHighlight()
 88 |     #         self.deSelectCorner()
 89 |     #     self.prevPoint = QPointF()
 90 |     #     self.repaint()
 91 | 
 92 | 
 93 |     def readDepth(self, point):
 94 |         u = point[0] / self.color_width * self.depth_width
 95 |         v = point[1] / self.color_height * self.depth_height
 96 |         return self.depth[int(round(v)), int(round(u))]
 97 | 
 98 |     def mousePressEvent(self, ev):
 99 |         #print(self.drawing(), pos)
100 |         point = self.transformPos(ev.pos())
101 |         if ev.button() == Qt.LeftButton:
102 |             self.scene.addCorner(point + self.patchOffsets, axisAligned=not self.shiftPressed)
103 |             pass
104 |         self.prevPoint = point
105 | 
106 |         #if ev.button() == Qt.RightButton:
107 |         #pos = self.transformPos(ev.pos(), moving=True)
108 |         #pass
109 |         self.repaint()
110 |         return
111 | 
112 |     def mouseMoveEvent(self, ev):
113 |         """Update line with last point and current coordinates."""
114 | 
115 |         if (Qt.RightButton & ev.buttons()):
116 |             point = self.transformPos(ev.pos())
117 |             self.patchOffsets -= (point - self.prevPoint).astype(np.int32)
118 |             sizes = np.array([self.scene.topdown.shape[1], self.scene.topdown.shape[0]])
119 |             self.patchOffsets = np.minimum(np.maximum(self.patchOffsets, 0), sizes - self.patchSizes)
120 |             #self.scene.reloadTopdown()
121 |             self.repaint()
122 |             #self.patchOffsets = np.minimum(np.maximum(self.patchOffsets, 0), self.patchSizes)
123 |             self.prevPoint = point
124 |             #self.loadTopdownImage()
125 |             return
126 | 
127 |         if (Qt.LeftButton & ev.buttons()):
128 |             point = self.transformPos(ev.pos())
129 |             self.scene.moveCorner(point - self.prevPoint)
130 |             self.repaint()
131 |             #self.patchOffsets = np.minimum(np.maximum(self.patchOffsets, 0), self.patchSizes)
132 |             self.prevPoint = point
133 |             #self.loadTopdownImage()
134 |             return
135 | 
136 |         return
137 | 
138 |     def mouseReleaseEvent(self, ev):
139 |         if self.ctrlPressed and self.shiftPressed and self.scene.selectedCornerIndex != -1:
140 |             point = self.transformPos(ev.pos())
141 |             self.scene.moveCorner(point, self.extrinsics_inv, self.intrinsics, self.imageIndex, recording=True)
142 |             pass
143 |         elif self.shiftPressed and self.scene.selectedCornerIndex != -1:
144 |             point = self.transformPos(ev.pos())
145 |             self.scene.moveCorner(point, self.extrinsics_inv, self.intrinsics, self.imageIndex, concave=True)
146 |             self.repaint()
147 |             pass
148 |         self.scene.selectedLayoutCorner = [-1, -1]
149 |         self.scene.selectedCornerIndex = -1
150 |         self.scene.selectedEdgeIndex = -1
151 | 
152 |         return
153 | 
154 |     def wheelEvent(self, ev):
155 |         if ev.delta() < 0:
156 |             self.topdownScale = max(self.topdownScale - 1, 1)
157 |         else:
158 |             self.topdownScale = self.topdownScale + 1
159 |             pass
160 |         #self.scene.reloadTopdown()
161 |         self.repaint()
162 |         return
163 | 
164 |     def handleDrawing(self, pos):
165 |         self.update()
166 | 
167 |     def selectCornerPoint(self, point):
168 |         """Select the first corner created which contains this point."""
169 |         self.deSelectCorner()
170 |         for corner in reversed(self.corners):
171 |             if corner.selectCorner(point, self.epsilon):
172 |                 self.selectCorner(corner)
173 |                 #self.calculateOffsets(corner, point)
174 |                 break
175 |             continue
176 |         return
177 | 
178 |     def paintEvent(self, event):
179 |         if (self.imageIndex == -1 or not self.image) and self.mode != 'layout':
180 |             return super(Canvas, self).paintEvent(event)
181 | 
182 |         p = self._painter
183 |         p.begin(self)
184 |         p.setRenderHint(QPainter.Antialiasing)
185 |         p.setRenderHint(QPainter.HighQualityAntialiasing)
186 |         p.setRenderHint(QPainter.SmoothPixmapTransform)
187 | 
188 |         self.scene.paint(p, self.patchOffsets, self.patchSizes, self.topdownOffset, self.topdownScale, self.offsetX, self.offsetY)
189 | 
190 |         p.end()
191 |         return
192 | 
193 |     def transformPos(self, point, moving=False):
194 |         """Convert from widget-logical coordinates to painter-logical coordinates."""
195 | 
196 |         return np.array([float(point.x() - self.offsetX), float(point.y() - self.offsetY)])
197 | 
198 | 
199 |     def closeEnough(self, p1, p2):
200 |         return distance(p1 - p2) < self.epsilon
201 | 
202 | 
203 |     def keyPressEvent(self, ev):
204 |         key = ev.key()
205 |         if (ev.modifiers() & Qt.ControlModifier):
206 |             self.ctrlPressed = True
207 |             if self.hiding:
208 |                 self.repaint()
209 |                 pass
210 |         else:
211 |             self.ctrlPressed = False
212 |             pass
213 |         if (ev.modifiers() & Qt.ShiftModifier):
214 |             self.shiftPressed = True
215 |             if self.hiding:
216 |                 self.repaint()
217 |                 pass
218 |         else:
219 |             self.shiftPressed = False
220 |             pass
221 | 
222 |         if key == Qt.Key_Escape:
223 |             #self.mode = 'moving'
224 |             self.scene.deleteSelected()
225 |             self.repaint()
226 |         if key == Qt.Key_Z:
227 |             self.scene.removeLast()
228 |             self.repaint()
229 |         elif key == Qt.Key_R:
230 |             if self.ctrlPressed:
231 |                 self.scene.reset('init')
232 |                 self.repaint()
233 |                 pass
234 |         elif key == Qt.Key_H:
235 |             #and Qt.ControlModifier == int(ev.modifiers()):
236 |             self.hiding = not self.hiding
237 |             self.repaint()
238 |         elif key == Qt.Key_A:
239 |             self.scene.finalize()
240 |         elif key == Qt.Key_Q:
241 |             if self.mode != 'layout':
242 |                 self.mode = 'point'
243 |                 pass
244 |         elif key == Qt.Key_S:
245 |             if self.ctrlPressed:
246 |                 print('save')
247 |                 self.scene.save()
248 |                 pass
249 |         elif key == Qt.Key_D:
250 |             self.setCurrentLabel(2)
251 |             self.setMode(False)
252 |         elif key == Qt.Key_F:
253 |             self.setCurrentLabel(3)
254 |             self.setMode(False)
255 |         elif key == Qt.Key_M:
256 |             self.writePLYFile()
257 |         elif key == Qt.Key_Right:
258 |             self.moveToNextImage()
259 |         elif key == Qt.Key_Left:
260 |             self.moveToPreviousImage()
261 |         elif key == Qt.Key_Down:
262 |             self.moveToNextImage(5)
263 |         elif key == Qt.Key_Up:
264 |             self.moveToPreviousImage(5)
265 |         elif key == Qt.Key_1:
266 |             self.moveToNextImage()
267 |             self.mode = 'move'
268 |             self.repaint()
269 |         elif key == Qt.Key_2:
270 |             self.showDensityImage()
271 |             self.mode = 'layout'
272 |             self.repaint()
273 |         elif key == Qt.Key_E:
274 |             if self.ctrlPressed:
275 |                 self.scene.exportPly()
276 |                 pass
277 |         elif key == Qt.Key_Space:
278 |             self.scene.finalize()
279 |             self.repaint()
280 |             pass
281 |         return
282 | 
283 |     def keyReleaseEvent(self, ev):
284 |         if self.hiding and self.ctrlPressed:
285 |             self.repaint()
286 |             pass
287 |         self.ctrlPressed = False
288 |         self.shiftPressed = False
289 |         return
290 | 
291 |     def setCurrentLabel(self, label):
292 |         self.currentLabel = label
293 |         return
294 | 
295 | 
296 | 
297 |     def loadCorners(self, corners):
298 |         self.corners = list(corners)
299 |         self.current = None
300 |         self.currentGroup = currentGroup
301 |         self.repaint()
302 | 
303 |     def onPoint(self, pos):
304 |         for corner in self.corners:
305 |             if corner.nearestVertex(pos, self.epsilon) is not None:
306 |                 return True
307 |             continue
308 |         return False
309 | 
310 | 
311 |     def loadScene(self, scenePath):
312 |         self.scene = Scene(scenePath)
313 | 
314 |         #self.scene.reloadTopdown()
315 |         self.repaint()
316 |         return
317 | 
318 |     def showDensityImage(self):
319 |         image = self.scene.getDensityImage(self.layout_width, self.layout_height)
320 |         self.image = QPixmap(QImage(image[:, :, ::-1].reshape(-1), self.layout_width, self.layout_height, self.layout_width * 3, QImage.Format_RGB888))
321 |         return
322 | 
323 |     def moveToNextImage(self, delta=1):
324 |         self.imageIndex = min(self.imageIndex + delta, len(self.imagePaths) - 1)
325 |         self.loadImage()
326 |         return
327 | 
328 |     def moveToPreviousImage(self, delta=1):
329 |         self.imageIndex = max(self.imageIndex - delta, 0)
330 |         self.loadImage()
331 |         return
332 | 
333 |     def loadImage(self):
334 |         image = cv2.imread(self.imagePaths[self.imageIndex])
335 |         self.image = QPixmap(QImage(image[:, :, ::-1].reshape(-1), self.color_width, self.color_height, self.color_width * 3, QImage.Format_RGB888))
336 | 
337 |         self.depth = cv2.imread(self.imagePaths[self.imageIndex].replace('color.jpg', 'depth.pgm'), -1).astype(np.float32) / 1000
338 | 
339 | 
340 |         self.extrinsics_inv = []
341 |         with open(self.imagePaths[self.imageIndex].replace('color.jpg', 'pose.txt'), 'r') as f:
342 |             for line in f:
343 |                 self.extrinsics_inv += [float(value) for value in line.strip().split(' ') if value.strip() != '']
344 |                 continue
345 |             pass
346 |         self.extrinsics_inv = np.array(self.extrinsics_inv).reshape((4, 4))
347 |         self.extrinsics = np.linalg.inv(self.extrinsics_inv)
348 |         self.repaint()
349 |         return
350 | 
351 |     # def loadTopdownImage(self):
352 |     #     topdown = self.scene.grabTopdown(self.patchOffsets, self.patchSizes)
353 |     #     #topdown = topdown[:, :, ::-1]
354 |     #     topdown = np.minimum((topdown - self.topdownOffset).astype(np.float32) / self.topdownScale * 255, 255).astype(np.uint8)
355 |     #     topdown = np.tile(np.expand_dims(topdown, axis=-1), [1, 1, 3])
356 |     #     self.topdownImage = QPixmap(QImage(topdown.reshape(-1), self.patchSizes[0], self.patchSizes[1], self.patchSizes[0] * 3, QImage.Format_RGB888))
357 |     #     self.repaint()
358 |     #     return
359 | 
360 |     def removeLastPoint(self):
361 |         self.corners = self.corners[:-1]
362 |         self.repaint()
363 |         return
364 | 
365 |     def sizeHint(self):
366 |         return self.minimumSizeHint()
367 | 
368 |     def minimumSizeHint(self):
369 |         if self.image:
370 |             return self.image.size()
371 |         return super(Canvas, self).minimumSizeHint()
372 | 


--------------------------------------------------------------------------------