class Window(QWidget):
def __init__(self, args):
super().__init__()
self.annotation_path = args.label_file
self.scale_factor = args.scale
self.annotaion_dict = dict()
self.img_list = list()
self.img_name = list()
img_list = os.listdir(args.img_dir)
for i in range(len(img_list)):
if img_list[i].endswith('jpg') or img_list[i].endswith('png'):
img_path = os.path.join(args.img_dir, img_list[i])
self.img_list.append(img_path)
self.img_name.append(img_list[i])
self.img_num = len(self.img_list)
self.img_id = 0
self.class_id = 1
if self.img_num == 0:
return
self._read_img()
self.region_img_view = ImageLabel(self)
self.region_img_view.setPixmap(
QPixmap(cvimg_to_qtimg(self.current_img_resize)))
self.region_img_view.setMouseTracking(True)
self.region_img_view.update_img(self.current_img_resize,
self.img_name[self.img_id])
self.button_width = 400 * self.scale_factor
self.button_height = 60 * self.scale_factor
self.pushButton1 = QPushButton(self)
self.pushButton1.setText("Next Img")
self.pushButton1.clicked.connect(self._down_func)
self.pushButton2 = QPushButton(self)
self.pushButton2.setText("Last Img")
self.pushButton2.clicked.connect(self._up_func)
self.pushButton3 = QPushButton(self)
self.pushButton3.setText("Apply")
self.pushButton3.clicked.connect(self._apply_func)
self.pushButton4 = QPushButton(self)
self.pushButton4.setText("Clean")
self.pushButton4.clicked.connect(self._clean_func)
self.pushButton5 = QPushButton(self)
self.pushButton5.setText("Output")
self.pushButton5.clicked.connect(self._output_func)
if args.mode == "multi":
self.class_text_height = 30 * self.scale_factor
self.class_text_width = 100 * self.scale_factor
self.class_text_label = QLabel(self)
self.class_text_label.setAlignment(Qt.AlignCenter)
self.class_text_label.setText('Class ID:')
self.setup_class_id = QLineEdit(self)
self.setup_class_id.setFocus()
self.setup_class_id.setFocusPolicy(Qt.ClickFocus)
self.setup_class_id.installEventFilter(self)
self.setup_class_id.editingFinished.connect(
self._setup_class_id_func)
else:
self.setup_class_id = QLabel(self)
self.setup_class_id.setFocus()
self.setup_class_id.setFocusPolicy(Qt.ClickFocus)
self.setup_class_id.installEventFilter(self)
self.setWindowTitle('Simple 2D Object Annotator')
self._init_window()
def eventFilter(self, source, event):
if (event.type() == QEvent.KeyPress and source is self.setup_class_id):
if event.key() == Qt.Key_Up:
self.pushButton2.click()
if event.key() == Qt.Key_Down:
self.pushButton1.click()
return super(Window, self).eventFilter(source, event)
def _setup_class_id_func(self):
setup_class_id_str = self.setup_class_id.text()
if setup_class_id_str != '':
try:
self.class_id = int(setup_class_id_str)
self.region_img_view.update_class(self.class_id)
except Exception as e:
print(e)
def _init_window(self):
window_width = self.img_width + self.button_width + 20
if args.mode == "multi":
other_height = 5 * (self.button_height + 10) + \
self.class_text_height + 10
self.setup_class_id.setGeometry(
QRect(self.img_width + 20 + self.class_text_width,
10 + (10 + self.button_height) * 5,
self.class_text_width, self.class_text_height))
self.class_text_label.setGeometry(
QRect(self.img_width + 10, 10 + (10 + self.button_height) * 5,
self.class_text_width, self.class_text_height))
else:
other_height = 5 * (self.button_height + 10)
window_height = max(self.img_height, other_height + 10)
self.setGeometry(100, 100, window_width, window_height)
self.region_img_view.setGeometry(
QRect(0, 0, self.img_width, self.img_height))
self.pushButton1.setGeometry(
QRect(self.img_width + 10, 10, self.button_width,
self.button_height))
self.pushButton2.setGeometry(
QRect(self.img_width + 10, 10 + (10 + self.button_height),
self.button_width, self.button_height))
self.pushButton3.setGeometry(
QRect(self.img_width + 10, 10 + (10 + self.button_height) * 2,
self.button_width, self.button_height))
self.pushButton4.setGeometry(
QRect(self.img_width + 10, 10 + (10 + self.button_height) * 3,
self.button_width, self.button_height))
self.pushButton5.setGeometry(
QRect(self.img_width + 10, 10 + (10 + self.button_height) * 4,
self.button_width, self.button_height))
def _read_img(self):
img_original = cv2.imread(self.img_list[self.img_id])
image_resize = cv2.resize(img_original,
None,
fx=self.scale_factor,
fy=self.scale_factor)
height, width, _ = image_resize.shape
self.current_img_resize = image_resize
self.current_img = img_original
self.img_height = height
self.img_width = width
def _reshape_bbox(self, bbox_list):
bbox_resize_list = list()
for item in bbox_list:
x = int(item['bbox'][0] / self.scale_factor)
y = int(item['bbox'][1] / self.scale_factor)
w = int(item['bbox'][2] / self.scale_factor)
h = int(item['bbox'][3] / self.scale_factor)
bbox_dict = {'bbox': [x, y, w, h], 'class': item['class']}
bbox_resize_list.append(bbox_dict)
return bbox_resize_list
def _down_func(self):
self.img_id += 1
if self.img_id == self.img_num:
self.img_id = self.img_num - 1
self._read_img()
self._init_window()
self.region_img_view.update_img(self.current_img_resize,
self.img_name[self.img_id])
self.region_img_view.setPixmap(
QPixmap(cvimg_to_qtimg(self.current_img_resize)))
self.region_img_view.update_class(self.class_id)
self.setup_class_id.setFocus()
def _up_func(self):
self.img_id -= 1
if self.img_id == -1:
self.img_id = 0
self._read_img()
self._init_window()
self.region_img_view.update_img(self.current_img_resize,
self.img_name[self.img_id])
self.region_img_view.setPixmap(
QPixmap(cvimg_to_qtimg(self.current_img_resize)))
self.region_img_view.update_class(self.class_id)
self.setup_class_id.setFocus()
def _apply_func(self):
image_local = self.current_img_resize.copy()
for item in self.region_img_view.bbox_list:
bbox = item['bbox']
class_id = item['class']
class_id_str = "class: " + str(class_id)
cv2.rectangle(image_local, (bbox[0], bbox[1]),
(bbox[0] + bbox[2], bbox[1] + bbox[3]), (0, 255, 0),
2)
cv2.putText(image_local, class_id_str, (bbox[0], bbox[1] + 15),
cv2.FONT_HERSHEY_TRIPLEX, 0.5, (0, 255, 0), 1)
self.region_img_view.setPixmap(QPixmap(cvimg_to_qtimg(image_local)))
bbox_resize_list = self._reshape_bbox(self.region_img_view.bbox_list)
self.annotaion_dict[self.img_name[self.img_id]] = bbox_resize_list
self.setup_class_id.setFocus()
def _clean_func(self):
self.region_img_view.clean_img()
self.region_img_view.setPixmap(
QPixmap(cvimg_to_qtimg(self.current_img_resize)))
self.annotaion_dict[self.img_name[self.img_id]] = []
self.setup_class_id.setFocus()
def _output_func(self):
json_str = json.dumps(self.annotaion_dict, indent=1)
with open(self.annotation_path, 'w') as file:
file.write(json_str)
self.setup_class_id.setFocus()
class DocumentWidget(QtGui.QWidget):
"""
DocumentWidget is the main component of MainWindow. It displays the PDF.
"""
def __init__(self, parent = None):
""" Initialize DocumentWidget. """
QtGui.QWidget.__init__(self, parent)
self.parent = parent
self.Document = None
self.CurrentPage = None
self.Image = None
self.highlights = None
self.num_pages = 0
self.page = 0
self.offset = 0
self.ImgLabel = ImageLabel(self)
self.ImgLabel.setAlignment(QtCore.Qt.AlignCenter)
self.ImgPixmap = QtGui.QPixmap()
self.scale = 1
def set_scale(self, event):
""" Sets the scale with which the document will be redered. """
self.scale = event
if self.Document is not None:
self.update_image()
def load_document(self, document):
"""
Load the document. Sets current page to the first page and the scale to 1.
"""
self.Document = popplerqt4.Poppler.Document.load(document)
self.Document.setRenderHint(self.Document.Antialiasing)
self.Document.setRenderHint(self.Document.TextAntialiasing)
self.Document.setRenderHint(self.Document.TextHinting)
self.CurrentPage = self.Document.page(self.page)
self.num_pages = self.Document.numPages()
self.set_scale(1)
def change_page(self, event):
""" Changes the page. """
if self.Document is not None:
self.page = ((event - self.offset -1) % self.num_pages)
#print 'Page %d.' % self.page
self.CurrentPage = self.Document.page(self.page)
self.update_image()
self.parent.ensureVisible(0, 0)
def fit_to_width_or_height(self, event):
"""
Changes the scale in such a way that it fits the width or the height of
the window.
width: event = 1
height: event = 2
"""
if self.Document is not None:
page_size = self.CurrentPage.pageSizeF()
if event == 1:
# 18 is window border, 4 is shadow
width = self.parent.rect().width() - 18 - 6
self.scale = 72.0*width/(DPI_X * page_size.width())
else:
# 2 is window border, 4 is shadow
height = self.parent.rect().height() - 2 - 6
self.scale = 72.0*height/(DPI_Y * page_size.height())
self.update_image()
def update_image(self):
""" Updates Image and ImgLabel. """
if self.Document is not None:
self.Image = self.CurrentPage.renderToImage(DPI_X*self.scale,
DPI_Y*self.scale)
self.ImgLabel.setPixmap(self.ImgPixmap.fromImage(self.Image))
background = QtGui.QImage(QtCore.QSize(self.Image.width() + 6,
self.Image.height() + 6),
self.Image.format())
painter = QtGui.QPainter()
painter.begin(background)
painter.fillRect(QtCore.QRect(0, 0, background.width(),
background.height()),
QtGui.QColor(60, 60, 60))
painter.drawImage(1, 1, self.Image)
if PLATFORM == 'Windows':
grey = QtGui.QColor(172, 168, 153)
else:
grey = QtGui.QColor(203, 201, 200)
painter.fillRect(QtCore.QRect(0, self.Image.height()+2, 6, 4), grey)
painter.fillRect(QtCore.QRect(self.Image.width()+ 2, 0, 4, 6), grey)
if self.highlights:
for rect, pg in self.highlights:
if pg == self.page:
(x, y) = self.ImgLabel.pt2px(QtCore.QSizeF(rect.x(),
rect.y()))
(w, h) = self.ImgLabel.pt2px(QtCore.QSizeF(rect.width(),
rect.height()))
painter.fillRect(QtCore.QRect(x + 1, y + 1, w, h),
QtGui.QColor(255, 255, 0, 100))
painter.end()
for note in self.ImgLabel.notes.values():
if note.page == self.page:
#print note.uid
x = note.pos.x()
x *= self.scale * DPI_X / 72.0
y = note.pos.y()
y *= self.scale * DPI_Y / 72.0
#print x,y
painter = QtGui.QPainter()
painter.begin(background)
painter.drawImage(x, y, self.ImgLabel.noteImage)
painter.end()
self.ImgLabel.setPixmap(QtGui.QPixmap.fromImage(background))
