Commit 7b3296ba authored by Immanuel Weber's avatar Immanuel Weber
Browse files

add support for center/dim voc export

parent 969dfe56
......@@ -297,7 +297,7 @@ class MainWindow(QMainWindow, WindowMixin):
checkable=True, enabled=False)
fitWidth = action('Fit &Width', self.setFitWidth,
'Ctrl+Shift+F', 'fit-width', u'Zoom follows window width',
checkable=True, enabled=False)
checkable=True, enabled=False)
# Group zoom controls into a list for easier toggling.
zoomActions = (self.zoomWidget, zoomIn, zoomOut,
zoomOrg, fitWindow, fitWidth)
......@@ -738,24 +738,22 @@ class MainWindow(QMainWindow, WindowMixin):
def loadLabels(self, shapes):
s = []
for label, points, line_color, fill_color, difficult in shapes:
for label, points, box, line_color, fill_color, difficult in shapes:
shape = Shape(label=label)
for x, y in points:
shape.addPoint(QPointF(x, y))
center = QPointF(*box[0])
width, height = box[1]
angle = box[2]
shape.addPoint(QPointF(center.x() - width / 2, center.y() - height / 2))
shape.addPoint(QPointF(center.x() + width / 2, center.y() - height / 2))
shape.addPoint(QPointF(center.x() + width / 2, center.y() + height / 2))
shape.addPoint(QPointF(center.x() - width / 2, center.y() + height / 2))
shape.applyRotationAngle(angle, center)
shape.updateState()
shape.difficult = difficult
shape.close()
s.append(shape)
if line_color:
shape.line_color = QColor(*line_color)
else:
shape.line_color = generateColorByText(label)
if fill_color:
shape.fill_color = QColor(*fill_color)
else:
shape.fill_color = generateColorByText(label)
shape.line_color = QColor(*line_color) if line_color else generateColorByText(label)
shape.fill_color = QColor(*fill_color) if fill_color else generateColorByText(label)
self.addLabel(shape)
self.canvas.loadShapes(s)
......@@ -767,12 +765,17 @@ class MainWindow(QMainWindow, WindowMixin):
self.labelFile.verified = self.canvas.verified
def format_shape(s):
s.updateState()
return dict(label=s.label,
line_color=s.line_color.getRgb(),
fill_color=s.fill_color.getRgb(),
points=[(p.x(), p.y()) for p in s.points],
# add chris
difficult = s.difficult)
difficult=s.difficult,
center=[s.center.x(), s.center.y()],
width=s.width,
height=s.height,
angle=s.currentAngle)
shapes = [format_shape(shape) for shape in self.canvas.shapes]
# Can add differrent annotation formats here
......
......@@ -638,7 +638,7 @@ class Canvas(QWidget):
for shift in shifts:
if shift is not None:
shape.points[index] += shift
shape.updateState()
# apply the new coordinates to the latent (unrotated) array
shape.applyRotationAngle(shape.currentAngle, shape_center, False)
......
......@@ -50,8 +50,13 @@ class LabelFile(object):
label = shape['label']
# Add Chris
difficult = int(shape['difficult'])
cx = shape['center'][0]
cy = shape['center'][1]
width = float(shape['width'])
height = float(shape['height'])
angle = float(shape['angle'])
bndbox = LabelFile.convertPoints2BndBox(points)
writer.addBndBox(bndbox[0], bndbox[1], bndbox[2], bndbox[3], label, difficult)
writer.addBndBox(bndbox[0], bndbox[1], bndbox[2], bndbox[3], label, difficult, cx, cy, width, height, angle)
writer.save(targetFile=filename)
return
......
......@@ -5,6 +5,7 @@ from xml.etree import ElementTree
from xml.etree.ElementTree import Element, SubElement
from lxml import etree
import codecs
import math
XML_EXT = '.xml'
ENCODE_METHOD = 'utf-8'
......@@ -74,10 +75,15 @@ class PascalVocWriter:
segmented.text = '0'
return top
def addBndBox(self, xmin, ymin, xmax, ymax, name, difficult):
def addBndBox(self, xmin, ymin, xmax, ymax, name, difficult, cx, cy, width, height, angle):
bndbox = {'xmin': xmin, 'ymin': ymin, 'xmax': xmax, 'ymax': ymax}
bndbox['name'] = name
bndbox['difficult'] = difficult
bndbox['cx'] = cx
bndbox['cy'] = cy
bndbox['width'] = width
bndbox['height'] = height
bndbox['angle'] = angle
self.boxlist.append(bndbox)
def appendObjects(self, top):
......@@ -109,6 +115,21 @@ class PascalVocWriter:
xmax.text = str(each_object['xmax'])
ymax = SubElement(bndbox, 'ymax')
ymax.text = str(each_object['ymax'])
realbox = SubElement(object_item, 'box')
cx = SubElement(realbox, 'cx')
cx.text = f"{each_object['cx']:.1f}"
cy = SubElement(realbox, 'cy')
cy.text = f"{each_object['cy']:.1f}"
width = SubElement(realbox, 'width')
width.text = f"{each_object['width']:.1f}"
height = SubElement(realbox, 'height')
height.text = f"{each_object['height']:.1f}"
angle = SubElement(realbox, 'angle')
angle.text = f"{each_object['angle']:.4f}"
angle_deg = SubElement(realbox, 'angle_deg')
angle_deg.text = f"{(each_object['angle'] / math.pi * 180):.1f}"
def save(self, targetFile=None):
root = self.genXML()
......@@ -141,13 +162,19 @@ class PascalVocReader:
def getShapes(self):
return self.shapes
def addShape(self, label, bndbox, difficult):
def addShape(self, label, bndbox, box, difficult):
xmin = int(bndbox.find('xmin').text)
ymin = int(bndbox.find('ymin').text)
xmax = int(bndbox.find('xmax').text)
ymax = int(bndbox.find('ymax').text)
points = [(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)]
self.shapes.append((label, points, None, None, difficult))
cx = float(box.find('cx').text)
cy = float(box.find('cy').text)
width = float(box.find('width').text)
height = float(box.find('height').text)
angle = float(box.find('angle').text)
box = ((cx, cy), (width, height), angle)
bbox_points = [(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)]
self.shapes.append((label, bbox_points, box, None, None, difficult))
def parseXML(self):
assert self.filepath.endswith(XML_EXT), "Unsupport file format"
......@@ -163,10 +190,11 @@ class PascalVocReader:
for object_iter in xmltree.findall('object'):
bndbox = object_iter.find("bndbox")
box = object_iter.find("box")
label = object_iter.find('name').text
# Add chris
difficult = False
if object_iter.find('difficult') is not None:
difficult = bool(int(object_iter.find('difficult').text))
self.addShape(label, bndbox, difficult)
self.addShape(label, bndbox, box, difficult)
return True
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment