Skip to content

Commit

Permalink
Update labelme to 5.4.0a0
Browse files Browse the repository at this point in the history
  • Loading branch information
@healthonrails
healthonrails committed Dec 4, 2023
1 parent bd0018c commit 99244bb
Show file tree
Hide file tree
Showing 5 changed files with 151 additions and 54 deletions.
1 change: 1 addition & 0 deletions annolid/configs/default_config.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@ canvas:
linestrip: false
polygonSAM: false
ai_polygon: false
ai_mask: false

# segment anything related parameters
sam:
Expand Down
14 changes: 14 additions & 0 deletions annolid/gui/app.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -620,6 +620,20 @@ def toggleDrawMode(self, edit=True, createMode="polygon"):
name=self._selectAiModelComboBox.currentText()
)
self._selectAiModelComboBox.setEnabled(True)
elif createMode == "ai_mask":
self.actions.createMode.setEnabled(True)
self.actions.createRectangleMode.setEnabled(True)
self.actions.createCircleMode.setEnabled(True)
self.actions.createLineMode.setEnabled(True)
self.actions.createPointMode.setEnabled(True)
self.actions.createLineStripMode.setEnabled(True)
self.actions.createAiPolygonMode.setEnabled(False)
self.canvas.initializeAiModel(
name=self._selectAiModelComboBox.currentText()
)
self._selectAiModelComboBox.setEnabled(True)
self.actions.createAiMaskMode.setEnabled(False)

else:
raise ValueError("Unsupported createMode: %s" % createMode)
self.actions.editMode.setEnabled(not edit)
Expand Down
116 changes: 74 additions & 42 deletions annolid/gui/shape.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from qtpy import QtGui
from labelme.logger import logger
import labelme.utils
import skimage.measure

# TODO(unknown):
# - [opt] Store paths instead of creating new ones at each paint.
Expand Down Expand Up @@ -53,6 +54,7 @@ def __init__(
flags=None,
group_id=None,
description=None,
mask=None,
):
self.label = label
self.group_id = group_id
Expand All @@ -67,6 +69,7 @@ def __init__(
self.flags = flags
self.description = description
self.other_data = {}
self.mask = mask

self._highlightIndex = None
self._highlightMode = self.NEAR_VERTEX
Expand All @@ -85,16 +88,17 @@ def __init__(

self.shape_type = shape_type

def setShapeRefined(self, points, point_labels, shape_type):
self._shape_raw = (self.points, self.point_labels, self.shape_type)
def setShapeRefined(self, shape_type, points, point_labels, mask=None):
self._shape_raw = (self.shape_type, self.points, self.point_labels)
self.shape_type = shape_type
self.points = points
self.point_labels = point_labels
self.shape_type = shape_type
self.mask = mask

def restoreShapeRaw(self):
if self._shape_raw is None:
return
self.points, self.point_labels, self.shape_type = self._shape_raw
self.shape_type, self.points, self.point_labels = self._shape_raw
self._shape_raw = None

@property
Expand All @@ -116,6 +120,7 @@ def shape_type(self, value):
"linestrip",
"multipoints",
"points",
"mask",
]:
raise ValueError("Unexpected shape_type: {}".format(value))
self._shape_type = value
Expand Down Expand Up @@ -188,26 +193,56 @@ def getRectFromLine(self, pt1, pt2):
return QtCore.QRectF(x1, y1, x2 - x1, y2 - y1)

def paint(self, painter):
if self.points:
color = (
self.select_line_color if self.selected else self.line_color
if self.mask is None and not self.points:
return

color = self.select_line_color if self.selected else self.line_color
pen = QtGui.QPen(color)
# Try using integer sizes for smoother drawing(?)
pen.setWidth(max(1, int(round(2.0 / self.scale))))
painter.setPen(pen)

if self.mask is not None:
image_to_draw = np.zeros(self.mask.shape + (4,), dtype=np.uint8)
fill_color = (
self.select_fill_color.getRgb()
if self.selected
else self.fill_color.getRgb()
)
image_to_draw[self.mask] = fill_color
qimage = QtGui.QImage.fromData(
labelme.utils.img_arr_to_data(image_to_draw)
)
painter.drawImage(
int(round(self.points[0].x())),
int(round(self.points[0].y())),
qimage,
)
pen = QtGui.QPen(color)
# Try using integer sizes for smoother drawing(?)
pen.setWidth(max(1, int(round(2.0 / self.scale))))
painter.setPen(pen)

line_path = QtGui.QPainterPath()
contours = skimage.measure.find_contours(
np.pad(self.mask, pad_width=1)
)
for contour in contours:
contour += [self.points[0].y(), self.points[0].x()]
line_path.moveTo(contour[0, 1], contour[0, 0])
for point in contour[1:]:
line_path.lineTo(point[1], point[0])
painter.drawPath(line_path)

if self.points:
line_path = QtGui.QPainterPath()
vrtx_path = QtGui.QPainterPath()
negative_vrtx_path = QtGui.QPainterPath()

if self.shape_type == "rectangle":
if self.shape_type in ["rectangle", "mask"]:
assert len(self.points) in [1, 2]
if len(self.points) == 2:
rectangle = self.getRectFromLine(*self.points)
line_path.addRect(rectangle)
for i in range(len(self.points)):
self.drawVertex(vrtx_path, i)
if self.shape_type == "rectangle":
for i in range(len(self.points)):
self.drawVertex(vrtx_path, i)
elif self.shape_type == "circle":
assert len(self.points) in [1, 2]
if len(self.points) == 2:
Expand All @@ -220,9 +255,8 @@ def paint(self, painter):
for i, p in enumerate(self.points):
line_path.lineTo(p)
self.drawVertex(vrtx_path, i)

elif self.shape_type == "points":
#assert len(self.points) == len(self.point_labels)
assert len(self.points) == len(self.point_labels)
for i, (p, l) in enumerate(
zip(self.points, self.point_labels)
):
Expand All @@ -242,40 +276,25 @@ def paint(self, painter):
self.drawVertex(vrtx_path, i)
if self.isClosed():
line_path.lineTo(self.points[0])
if self.shape_type == "line":

label_x, label_y = self.points[0].x(
), self.points[0].y()
if len(self.points) > 1:
dist = labelme.utils.distance(
self.points[0] - self.points[1])
painter.setFont(QtGui.QFont(
"Arial", 3 * int(self.point_size/self.scale)))
painter.drawText(
int(label_x - self.point_size),
int(label_y - self.point_size),
f"length:{round(dist,2)}pixels")

if self.label:
font = QtGui.QFont(
"Arial", int(3 * self.point_size/self.scale))
painter.setFont(font)
label_x, label_y = self.find_polygon_center(
self.points)
painter.drawText(
int(label_x)+1, int(label_y), str(self.label))

painter.drawPath(line_path)
painter.drawPath(vrtx_path)
painter.fillPath(vrtx_path, self._vertex_fill_color)
if self.fill:
if vrtx_path.length() > 0:
painter.drawPath(vrtx_path)
painter.fillPath(vrtx_path, self._vertex_fill_color)
if self.fill and self.mask is None:
color = (
self.select_fill_color
if self.selected
else self.fill_color
)
painter.fillPath(line_path, color)

pen.setColor(QtGui.QColor(255, 0, 0, 255))
painter.setPen(pen)
painter.drawPath(negative_vrtx_path)
painter.fillPath(negative_vrtx_path, QtGui.QColor(255, 0, 0, 255))


def drawVertex(self, path, i):
d = self.point_size / self.scale
shape = self.point_type
Expand Down Expand Up @@ -316,6 +335,19 @@ def nearestEdge(self, point, epsilon):
return post_i

def containsPoint(self, point):
if self.mask is not None:
y = np.clip(
int(round(point.y() - self.points[0].y())),
0,
self.mask.shape[0] - 1,
)
x = np.clip(
int(round(point.x() - self.points[0].x())),
0,
self.mask.shape[1] - 1,
)
return self.mask[y, x]

if not self.makePath():
return False
return self.makePath().contains(point)
Expand All @@ -331,7 +363,7 @@ def getCircleRectFromLine(self, line):
return rectangle

def makePath(self):
if self.shape_type == "rectangle":
if self.shape_type in ["rectangle", "mask"]:
path = QtGui.QPainterPath()
if len(self.points) == 2:
rectangle = self.getRectFromLine(*self.points)
Expand Down
Loading

0 comments on commit 99244bb

Please sign in to comment.