代码收藏家技术教程 2024-11-13

【标注工具】labelme5.0.1自定义功能——添加水平线标注及显示信息

一、背景

随着项目的逐步开展，数据标注的工作量越来越大，为了提高标注效率和标注准确性，需要对labelme的功能进行扩充，更好地适应标注需求，同时后期也能将数据标注工作顺利交付出去。

本次标注需要一个rectangle，一条水平line，一个point，并计算三者之间的相对位置关系，将计算结果显示到界面上，有助于精细化调整，因此需要添加水平线标注及显示信息两个功能。

二、具体工作

1. 添加水平线HLine标注

全局搜索“line”，仿照“line”的生成方式添加“hline”。需要修改的文件内容如下：

（1）utils/shape.py

Line36 后面添加代码

# 新增shape_type
elif shape_type == "hline":
    assert len(xy) == 2, "Shape of shape_type=hline must have 2 points"
    draw.line(xy=xy, fill=1, width=line_width)

（2）utils/canvas.py

Line103 createMode()中添加代码

def createMode(self, value):
        if value not in [
            "polygon",
            "rectangle",
            "circle",
            "line",
            "hline",    # 添加hline类型
            "point",
            "linestrip",
        ]:

Line236 后面添加代码段

# 添加水平线标注方法
elif self.createMode == "hline":
    pos2 = pos
    pos2.setY(self.current[0].y())    #将第二个点的y坐标替换为第一个点的y坐标
    self.line.points = [self.current[0], pos2]
    self.line.close()

Line362 mousePressEvent()中添加

elif self.createMode in ["rectangle", "circle", "line", "hline"]:    # 添加hline
    assert len(self.current.points) == 1
    self.current.points = self.line.points

Line837 undoLastLine()中添加

elif self.createMode in ["rectangle", "line", "circle", "hline"]:
    self.current.points = self.current.points[0:1]

（3）app.py

Line351 后面添加代码段

# 创建水平线方法
createHLineMode = action(
    self.tr("Create HLine"),
    lambda: self.toggleDrawMode(False, createMode="hline"),
    shortcuts["create_hline"],
    "objects",
    self.tr("Start drawing hlines"),
    enabled=False,
)

Line601 后添加代码

createHLineMode=createHLineMode

Line635、653 后添加代码

createHLineMode

Line850 后添加代码

self.actions.createHLineMode

Line881、958、966、973、980、987、994、1001后添加代码

self.actions.createLineMode.setEnabled(True)

Line982 后添加代码段

elif createMode == "hline":
    self.actions.createMode.setEnabled(True)
    self.actions.createRectangleMode.setEnabled(True)
    self.actions.createCircleMode.setEnabled(True)
    self.actions.createLineMode.setEnabled(True)
    self.actions.createHLineMode.setEnabled(False)
    self.actions.createPointMode.setEnabled(True)
    self.actions.createLineStripMode.setEnabled(True)

（4）shape.py

Line93 后面添加

if value not in [
    "polygon",
    "rectangle",
    "point",
    "line",
    "hline",    # 添加hline
    "circle",
    "linestrip",
]:

2. 添加信息显示

本次标注需要将点与线段间的距离值显示到标注界面上，随着点的移动实时更新距离值，主要修改canvas.py，创建rectangle、hline和point后提取相应点的 $y$ 坐标，在编辑形状的过程中，移动point的位置后更新 $y$ 坐标，计算新的距离值。

（1）坐标值记录

创建变量记录坐标值，class canvas()的__init__()

self.bottom_y = 0    # rectangle(x1,y1,x2,y2)的y2
self.top_y = 0    # hline(x1,y,x,2y)的y
self.move_p_y = 0    # point(x,y)的y
self.ht = 0.        # 距离值
self.ht_tag = False    # 是否开始计算距离值

创建rectangle后记录其 $y_2$ 的值

elif self.createMode == "rectangle":
    self.line.points = [self.current[0], pos]
    self.bottom_y = pos.y()    # 记录y2
    self.line.close()

创建hline后记录其 $y$ 值

elif self.createMode == "hline":
    pos2 = pos
    pos2.setY(self.current[0].y())
    self.line.points = [self.current[0], pos2]
    self.top_y = self.current[0].y()    # 记录y
    self.line.close()

创建point后记录其 $y$ 值

elif self.createMode == "point":
    self.line.points = [self.current[0]]
    self.move_p_y = self.current[0].y()    # 记录y
    self.ht = min(1., max(0., (self.bottom_y - self.move_p_y)) / (self.bottom_y -self.top_y))    # 计算距离值
    self.ht_tag = True    # 完成point标注，可以显示
    self.line.close()

（2）坐标更新

在鼠标点击坐标点或编辑形状时，也要实时记录，mousePressEvent()

if ev.button() == QtCore.Qt.LeftButton:
    if self.drawing():
        if self.current:
            # Add point to existing shape.
            if self.createMode == "polygon":
                self.current.addPoint(self.line[1])
                self.line[0] = self.current[-1]
                if self.current.isClosed():
                    self.finalise()
            elif self.createMode in ["rectangle", "circle", "line", "hline"]:
                assert len(self.current.points) == 1
                self.current.points = self.line.points
                if self.createMode == "rectangle":
                    self.bottom_y = self.current.points[1].y()    # 记录rec y2
                elif self.createMode == "hline":
                    self.top_y = self.current.points[0].y()    # 记录hline y
                self.finalise()
            elif self.createMode == "linestrip":
                self.current.addPoint(self.line[1])
                self.line[0] = self.current[-1]
                if int(ev.modifiers()) == QtCore.Qt.ControlModifier:
                    self.finalise()
        elif not self.outOfPixmap(pos):
            # Create new shape.
            self.current = Shape(shape_type=self.createMode)
            self.current.addPoint(pos)
            if self.createMode == "point":
                self.move_p_y = self.current.points[0].y()    # 记录point y
                self.ht = min(1., max(0., (self.bottom_y - self.move_p_y)) / (self.bottom_y - self.top_y))    # 更新距离值
                self.ht_tag = True
                self.finalise()
            else:
                if self.createMode == "circle":
                    self.current.shape_type = "circle"
                self.line.points = [pos, pos]
                self.setHiding()
                self.drawingPolygon.emit(True)
                self.update()
    elif self.editing():
        if self.selectedEdge():
            self.addPointToEdge()
        elif (
            self.selectedVertex()
            and int(ev.modifiers()) == QtCore.Qt.ShiftModifier
        ):
            # Delete point if: left-click + SHIFT on a point
            self.removeSelectedPoint()

        group_mode = int(ev.modifiers()) == QtCore.Qt.ControlModifier
        self.selectShapePoint(pos, multiple_selection_mode=group_mode)
        self.prevPoint = pos
        # 编辑形状时，根据调整的point的位置更新距离值
        if self.createMode == 'point':
            self.move_p_y = int(pos.y())
            self.ht = min(1., max(0., (self.bottom_y - self.move_p_y)) / (self.bottom_y - self.top_y))
            self.ht_tag = True

        self.repaint()

更新坐标，mouseReleaseEvent()

（2024-08-13修改，修复删除水平线时报错问题）

if self.movingShape and self.hShape:
    index = self.shapes.index(self.hShape)
    if (
        self.shapesBackups[-1][index].points
        != self.shapes[index].points
    ):
        if self.hShape.shape_type == 'rectangle':
            self.bottom_y = self.shapes[index].points[1].y()
        if self.hShape.shape_type == 'hline':
            # 修改，删除点时不更新坐标
            if len(self.shapes[index].points) > 1:
                self.shapes[index].points[1].setY(self.shapesBackups[-1][index].points[0].y())
            self.top_y = self.shapes[index].points[0].y()
        if self.hShape.shape_type == 'point':
            self.move_p_y = self.shapes[index].points[0].y()
            self.ht = min(1., max(0., (self.bottom_y - self.move_p_y)) / (self.bottom_y - self.top_y))
            self.ht_tag = True
        self.storeShapes()
        self.shapeMoved.emit()

    self.movingShape = False

（3）界面绘制

主要参考[1]

for shape in self.shapes:
    if (shape.selected or not self._hideBackround) and self.isVisible(
        shape
    ):
        shape.fill = shape.selected or shape == self.hShape
        shape.paint(p)
        p.setPen(QColor(0 ,255, 0))
        p.setFont(QFont('微软雅黑', 20))
        if self.ht_tag:
            labelText = f'当前料位高度：{self.ht}'
            p.drawText(QPointF(10, 40), labelText)