|
@@ -2,7 +2,8 @@ import imgviz
|
|
|
from qtpy import QtCore
|
|
|
from qtpy import QtGui
|
|
|
from qtpy import QtWidgets
|
|
|
-
|
|
|
+import cv2
|
|
|
+import numpy as np
|
|
|
import labelme.ai
|
|
|
import labelme.utils
|
|
|
from labelme import QT5
|
|
@@ -10,6 +11,8 @@ from labelme.logger import logger
|
|
|
from labelme.shape import Shape
|
|
|
import numpy as np
|
|
|
import cv2
|
|
|
+from labelme.utils import load_barcode_dict
|
|
|
+from labelme.widgets.label_dialog import LabelDialog
|
|
|
# TODO(unknown):
|
|
|
# - [maybe] Find optimal epsilon value.
|
|
|
|
|
@@ -22,7 +25,6 @@ CURSOR_GRAB = QtCore.Qt.OpenHandCursor
|
|
|
|
|
|
MOVE_SPEED = 5.0
|
|
|
|
|
|
-
|
|
|
class Canvas(QtWidgets.QWidget):
|
|
|
zoomRequest = QtCore.Signal(int, QtCore.QPoint)
|
|
|
scrollRequest = QtCore.Signal(int, int)
|
|
@@ -44,7 +46,10 @@ class Canvas(QtWidgets.QWidget):
|
|
|
self.epsilon = kwargs.pop("epsilon", 10.0)
|
|
|
self.detection_model_path=kwargs.pop("detection_model_path",None)
|
|
|
self.segmentation_model_path=kwargs.pop("segmentation_model_path",None)
|
|
|
+ self.decoding_model_path=kwargs.pop("decoding_model_path", None)
|
|
|
self.double_click = kwargs.pop("double_click", "close")
|
|
|
+ self._characters = load_barcode_dict() # Load dictionary once
|
|
|
+ logger.debug(f"Loaded barcode dictionary with {len(self._characters)} characters.")
|
|
|
if self.double_click not in [None, "close"]:
|
|
|
raise ValueError(
|
|
|
"Unexpected value for double_click event: {}".format(self.double_click)
|
|
@@ -99,6 +104,7 @@ class Canvas(QtWidgets.QWidget):
|
|
|
self.draw_pred=False
|
|
|
self.pred_bbox_points=None
|
|
|
self.current_bbox_point=None
|
|
|
+ self.decoded_barcode = None
|
|
|
# Menus:
|
|
|
# 0: right-click without selection and dragging of shapes
|
|
|
# 1: right-click with selection and dragging of shapes
|
|
@@ -110,6 +116,7 @@ class Canvas(QtWidgets.QWidget):
|
|
|
self._ai_model = None
|
|
|
self._detection_model = None
|
|
|
self._segmentation_model = None
|
|
|
+ self._decoding_model = None
|
|
|
def fillDrawing(self):
|
|
|
return self._fill_drawing
|
|
|
|
|
@@ -135,18 +142,25 @@ class Canvas(QtWidgets.QWidget):
|
|
|
raise ValueError("Unsupported createMode: %s" % value)
|
|
|
self._createMode = value
|
|
|
|
|
|
- def initializeBarcodeModel(self, detection_model_path, segmentation_model_path=None):
|
|
|
+ def initializeBarcodeModel(self, detection_model_path, segmentation_model_path=None, decoding_model_path=None):
|
|
|
if not detection_model_path:
|
|
|
raise ValueError("Detection model path is required.")
|
|
|
|
|
|
logger.debug("Initializing only detection model: %r" % "BarcodePredictModel")
|
|
|
- self._detection_model = labelme.ai.BarcodePredictModel(detection_model_path)
|
|
|
+ self._detection_model = labelme.ai.BarcodeDetectModel(detection_model_path)
|
|
|
|
|
|
if segmentation_model_path:
|
|
|
logger.debug("Initializing barcode detection & Segmentation model: %r" % "BarcodePredictModel")
|
|
|
- self._segmentation_model = labelme.ai.BarcodePredictModel(
|
|
|
+ self._segmentation_model = labelme.ai.BarcodeDetectModel(
|
|
|
detection_model_path, segmentation_model_path
|
|
|
)
|
|
|
+
|
|
|
+ if decoding_model_path:
|
|
|
+ logger.debug("Initializing barcode detection, Segmentation model, Decoding_model: %r" % "BarcodePredictModel")
|
|
|
+ self._decoding_model = labelme.ai.BarcodeDecodeModel(
|
|
|
+ decoding_model_path
|
|
|
+ )
|
|
|
+
|
|
|
if self.pixmap is None:
|
|
|
logger.warning("Pixmap is not set yet")
|
|
|
return
|
|
@@ -198,7 +212,7 @@ class Canvas(QtWidgets.QWidget):
|
|
|
# and app.py::loadShapes and our own Canvas::loadShapes function.
|
|
|
if not self.isShapeRestorable:
|
|
|
return
|
|
|
- print(f"shape is restorable")
|
|
|
+ # print(f"shape is restorable")
|
|
|
self.shapesBackups.pop() # latest
|
|
|
|
|
|
# The application will eventually call Canvas.loadShapes which will
|
|
@@ -273,6 +287,7 @@ class Canvas(QtWidgets.QWidget):
|
|
|
|
|
|
# Polygon drawing.
|
|
|
if self.drawing():
|
|
|
+
|
|
|
if self.createMode in ["ai_polygon", "ai_mask"]:
|
|
|
self.line.shape_type = "points"
|
|
|
else:
|
|
@@ -323,6 +338,7 @@ class Canvas(QtWidgets.QWidget):
|
|
|
self.line.points = [self.current[0]]
|
|
|
self.line.point_labels = [1]
|
|
|
self.line.close()
|
|
|
+ # print("self.line.points", self.line.points)
|
|
|
assert len(self.line.points) == len(self.line.point_labels)
|
|
|
self.repaint()
|
|
|
self.current.highlightClear()
|
|
@@ -846,6 +862,18 @@ class Canvas(QtWidgets.QWidget):
|
|
|
def finalise(self):
|
|
|
if(self.current is None):
|
|
|
return
|
|
|
+ # If in manual mode and points are finalized
|
|
|
+ if self.createMode in ["polygon", "rectangle"] and len(self.current.points) > 0:
|
|
|
+ if self._detection_model is None:
|
|
|
+ logger.info(f"Initializing AI model")
|
|
|
+ self.initializeBarcodeModel(self.detection_model_path, self.segmentation_model_path, self.decoding_model_path)
|
|
|
+ # Extract the points from the manually drawn shape
|
|
|
+ manual_points = [[point.x(), point.y()] for point in self.current.points]
|
|
|
+ print(manual_points)
|
|
|
+ logger.debug(f"Manually drawn points: {manual_points}")
|
|
|
+ self.detect_and_segment(manual_points=manual_points)
|
|
|
+ self.mode = self.EDIT # Or any appropriate mode to disable drawing behavior
|
|
|
+ return
|
|
|
if self.createMode == "ai_polygon":
|
|
|
# convert points to polygon by an AI model
|
|
|
assert self.current.shape_type == "points"
|
|
@@ -853,6 +881,7 @@ class Canvas(QtWidgets.QWidget):
|
|
|
points=[[point.x(), point.y()] for point in self.current.points],
|
|
|
point_labels=self.current.point_labels,
|
|
|
)
|
|
|
+
|
|
|
self.current.setShapeRefined(
|
|
|
points=[QtCore.QPointF(point[0], point[1]) for point in points],
|
|
|
point_labels=[1] * len(points),
|
|
@@ -882,11 +911,15 @@ class Canvas(QtWidgets.QWidget):
|
|
|
points=[QtCore.QPointF(point[0], point[1]) for point in bbox_point],
|
|
|
point_labels=[1]*len(bbox_point)
|
|
|
)
|
|
|
+ drawing_shape.description = self.decoded_barcode
|
|
|
drawing_shape.close()
|
|
|
self.shapes.append(drawing_shape)
|
|
|
self.storeShapes()
|
|
|
self.update()
|
|
|
self.newShape.emit()
|
|
|
+ # self.pred_bbox_points = None
|
|
|
+ # self.draw_pred = False
|
|
|
+ # self.detect_and_segment() # Process the next detection
|
|
|
current_copy.close()
|
|
|
current_copy=None
|
|
|
if(self.current):
|
|
@@ -1032,7 +1065,7 @@ class Canvas(QtWidgets.QWidget):
|
|
|
)
|
|
|
|
|
|
if self._detection_model:
|
|
|
- if self._segmentation_model is None:
|
|
|
+ if self.segmentation_model_path is None:
|
|
|
logger.info(f"Performing detection only.")
|
|
|
# Get prediction from model
|
|
|
self.pred_bbox_points = self._detection_model.predict_polygon_from_points()
|
|
@@ -1043,8 +1076,9 @@ class Canvas(QtWidgets.QWidget):
|
|
|
else:
|
|
|
print("No bounding boxes detected.")
|
|
|
else:
|
|
|
- logger.info(f"Performing detection and segmentation.")
|
|
|
- self.detect_and_segment()
|
|
|
+ # logger.info(f"Performing detection, segmentation and decoding.")
|
|
|
+ self.initializeBarcodeModel(self.detection_model_path, self.segmentation_model_path, self.decoding_model_path)
|
|
|
+ self.detect_and_segment(manual_points=None)
|
|
|
elif self.editing():
|
|
|
if key == QtCore.Qt.Key_Up:
|
|
|
self.moveByKeyboard(QtCore.QPointF(0.0, -MOVE_SPEED))
|
|
@@ -1055,13 +1089,57 @@ class Canvas(QtWidgets.QWidget):
|
|
|
elif key == QtCore.Qt.Key_Right:
|
|
|
self.moveByKeyboard(QtCore.QPointF(MOVE_SPEED, 0.0))
|
|
|
|
|
|
-
|
|
|
- def scale_points(self, approx, mask_shape, cropped_shape, x_min, y_min):
|
|
|
+ def scale_points(self, points, mask_shape, cropped_shape, x_min, y_min):
|
|
|
+ """
|
|
|
+ Scale the points from the segmentation mask space back to the original cropped image coordinates.
|
|
|
+
|
|
|
+ Args:
|
|
|
+ points (np.ndarray): Points to scale (Nx2 array or list of lists).
|
|
|
+ mask_shape (tuple): Shape of the segmentation mask (height, width).
|
|
|
+ cropped_shape (tuple): Shape of the cropped image in the original image (height, width).
|
|
|
+ x_min (int): Minimum x-coordinate of the cropped region in the original image.
|
|
|
+ y_min (int): Minimum y-coordinate of the cropped region in the original image.
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ scaled_points (list): List of scaled points as [[x1, y1], [x2, y2], ...].
|
|
|
+ """
|
|
|
+ # Compute scaling factors
|
|
|
scale_x = cropped_shape[1] / mask_shape[1] # Scale factor for x-axis
|
|
|
scale_y = cropped_shape[0] / mask_shape[0] # Scale factor for y-axis
|
|
|
- return [[int(pt[0][0] * scale_x) + x_min, int(pt[0][1] * scale_y) + y_min] for pt in approx]
|
|
|
+
|
|
|
+ # Scale and translate the points
|
|
|
+ scaled_points = []
|
|
|
+ for pt in points:
|
|
|
+ scaled_x = int(pt[0] * scale_x) + x_min
|
|
|
+ scaled_y = int(pt[1] * scale_y) + y_min
|
|
|
+ scaled_points.append([scaled_x, scaled_y])
|
|
|
+ # logger.debug(f"scaled points: {scaled_points}")
|
|
|
+ return scaled_points
|
|
|
|
|
|
- def detect_and_segment(self):
|
|
|
+ def expand_bbox(self, x_min, y_min, x_max, y_max, factor=1.3):
|
|
|
+ """
|
|
|
+ Expands the bounding box by a given factor.
|
|
|
+ """
|
|
|
+ center_x = (x_max + x_min) / 2
|
|
|
+ center_y = (y_max + y_min) / 2
|
|
|
+ width = (x_max - x_min) * factor
|
|
|
+ height = (y_max - y_min) * factor
|
|
|
+
|
|
|
+ new_x_min = int(center_x - width / 2)
|
|
|
+ new_y_min = int(center_y - height / 2)
|
|
|
+ new_x_max = int(center_x + width / 2)
|
|
|
+ new_y_max = int(center_y + height / 2)
|
|
|
+
|
|
|
+ # Ensure the coordinates do not go out of image bounds
|
|
|
+ new_x_min = max(0, new_x_min)
|
|
|
+ new_y_min = max(0, new_y_min)
|
|
|
+ new_x_max = min(self.pixmap.width() - 1, new_x_max)
|
|
|
+ new_y_max = min(self.pixmap.height() - 1, new_y_max)
|
|
|
+
|
|
|
+ return new_x_min, new_y_min, new_x_max, new_y_max
|
|
|
+
|
|
|
+
|
|
|
+ def detect_and_segment(self, manual_points=None):
|
|
|
"""
|
|
|
Perform detection and segmentation (if both models are available).
|
|
|
"""
|
|
@@ -1069,10 +1147,11 @@ class Canvas(QtWidgets.QWidget):
|
|
|
self.current = Shape(
|
|
|
shape_type="points" if self.createMode in ["ai_polygon", "ai_mask"] else self.createMode
|
|
|
)
|
|
|
+ print("self.current",self.current)
|
|
|
|
|
|
# Step 1: detection bounding box points
|
|
|
- detection_results = self._detection_model.predict_polygon_from_points()
|
|
|
-
|
|
|
+ detection_results = [manual_points] if manual_points else self._detection_model.predict_polygon_from_points()
|
|
|
+ print(detection_results)
|
|
|
if not detection_results or len(detection_results) == 0:
|
|
|
logger.warning("No detection found")
|
|
|
return
|
|
@@ -1099,7 +1178,7 @@ class Canvas(QtWidgets.QWidget):
|
|
|
continue
|
|
|
|
|
|
# Converting bounding box values to integers for cropping
|
|
|
- x_min, y_min, x_max, y_max = int(x_min), int(y_min), int(x_max), int(y_max)
|
|
|
+ x_min, y_min, x_max, y_max = self.expand_bbox(x_min, y_min, x_max, y_max)
|
|
|
|
|
|
# Step 3: Cropping image based on detection output
|
|
|
try:
|
|
@@ -1153,14 +1232,20 @@ class Canvas(QtWidgets.QWidget):
|
|
|
# Normalize the mask to 0 and 255 and convert to uint8
|
|
|
mask = (mask * 255).astype(np.uint8)
|
|
|
|
|
|
- logger.debug(f"Converted mask shape for detection {detection_idx + 1}: {mask.shape}, dtype: {mask.dtype}")
|
|
|
- cv2.imwrite(f"segmentation_mask_{detection_idx + 1}.png", mask)
|
|
|
+ # logger.debug(f"Converted mask shape for detection {detection_idx + 1}: {mask.shape}, dtype: {mask.dtype}")
|
|
|
+ # cv2.imwrite(f"segmentation_mask_{detection_idx + 1}.png", mask)
|
|
|
+
|
|
|
+ # Scale mask to original cropped image size
|
|
|
+ scaled_mask = cv2.resize(mask, (orig_cropped_shape[1], orig_cropped_shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
|
+
|
|
|
if rotated:
|
|
|
cropped_image = cv2.rotate(cropped_image, cv2.ROTATE_90_COUNTERCLOCKWISE)
|
|
|
- mask = cv2.rotate(mask, cv2.ROTATE_90_COUNTERCLOCKWISE)
|
|
|
+ mask = cv2.rotate(scaled_mask, cv2.ROTATE_90_COUNTERCLOCKWISE)
|
|
|
rotated_cropped_shape = cropped_image.shape[:2]
|
|
|
+ else:
|
|
|
+ mask = scaled_mask
|
|
|
|
|
|
- # cv2.imwrite(f"segmentation_mask_{detection_idx + 1}.png", mask)
|
|
|
+ cv2.imwrite(f"scaled_segmentation_mask_{detection_idx + 1}.png", mask)
|
|
|
logger.debug(f"Saved segmentation mask for detection {detection_idx + 1}.")
|
|
|
|
|
|
# Step 7: Find contours
|
|
@@ -1168,31 +1253,53 @@ class Canvas(QtWidgets.QWidget):
|
|
|
logger.debug(f"Found {len(contours)} contours in the mask for detection {detection_idx + 1}.")
|
|
|
|
|
|
if len(contours) > 0:
|
|
|
+ # Get the largest contour
|
|
|
largest_contour = max(contours, key=cv2.contourArea)
|
|
|
-
|
|
|
- # Step 8: Approximate a polygon with exactly 4 points (quadrilateral)
|
|
|
- epsilon = 0.02 * cv2.arcLength(largest_contour, True) # epsilon for precision
|
|
|
- approx = cv2.approxPolyDP(largest_contour, epsilon, True)
|
|
|
-
|
|
|
- # If the approximation doesn't result in 4 points, force it
|
|
|
- if len(approx) != 4:
|
|
|
- # Using boundingRect as fallback in case of insufficient points
|
|
|
- print("log here")
|
|
|
- x, y, w, h = cv2.boundingRect(largest_contour)
|
|
|
- point_xy = [
|
|
|
- [x + x_min, y + y_min], # Top-left
|
|
|
- [x + w + x_min, y + y_min], # Top-right
|
|
|
- [x + w + x_min, y + h + y_min], # Bottom-right
|
|
|
- [x + x_min, y + h + y_min] # Bottom-left
|
|
|
- ]
|
|
|
+
|
|
|
+ # Get bounding and rotated rectangles
|
|
|
+ rot_rect = cv2.minAreaRect(largest_contour)
|
|
|
+ box_points = cv2.boxPoints(rot_rect)
|
|
|
+ box_points = np.int32(box_points)
|
|
|
+
|
|
|
+ # Log the rotated rectangle points
|
|
|
+ # logger.debug(f"Rotated corners (box_points) for detection {detection_idx + 1}: {box_points}")
|
|
|
+
|
|
|
+ # Step 1: Draw the rotated rectangle on the cropped image
|
|
|
+ cropped_with_rects = cropped_image.copy()
|
|
|
+ cv2.drawContours(cropped_with_rects, [box_points], -1, (0, 255, 0), 2)
|
|
|
+ cv2.imwrite(f"cropped_with_rects_{detection_idx + 1}.png", cropped_with_rects)
|
|
|
+ logger.debug(f"Saved cropped image with rectangles for detection {detection_idx + 1}.")
|
|
|
+
|
|
|
+ if rotated:
|
|
|
+ # Scale points from mask space to the original image space with rotation
|
|
|
+ scaled_points = self.scale_points(box_points, mask.shape, rotated_cropped_shape, x_min, y_min)
|
|
|
else:
|
|
|
- if rotated:
|
|
|
- point_xy = self.scale_points(approx, mask.shape, rotated_cropped_shape, x_min, y_min)
|
|
|
- else:
|
|
|
- point_xy = self.scale_points(approx, mask.shape, orig_cropped_shape, x_min, y_min)
|
|
|
- logger.debug(f"Generated 4 corner points for the polygon for detection {detection_idx + 1}: {point_xy}")
|
|
|
- self.pred_bbox_points = [point_xy]
|
|
|
- logger.debug(f"Predicted Bounding Box Points for detection {detection_idx + 1}: {self.pred_bbox_points}")
|
|
|
+ # Scale points without rotation
|
|
|
+ scaled_points = self.scale_points(box_points, mask.shape, orig_cropped_shape, x_min, y_min)
|
|
|
+
|
|
|
+ self.pred_bbox_points = [scaled_points]
|
|
|
+ logger.debug(f"Predicted Bounding Box Points for barcode {detection_idx + 1}: {self.pred_bbox_points}")
|
|
|
+
|
|
|
+ if self.pred_bbox_points:
|
|
|
+ for bbox_points in self.pred_bbox_points:
|
|
|
+ try:
|
|
|
+ original_image = labelme.utils.img_qt_to_arr(self.pixmap.toImage())
|
|
|
+ # Call the decoding function
|
|
|
+ decoded_output = self._decoding_model.decode_from_points(
|
|
|
+ points=bbox_points,
|
|
|
+ detection_idx=detection_idx,
|
|
|
+ original_image = original_image
|
|
|
+ )
|
|
|
+ logger.debug(f"Decoded output for detection {detection_idx + 1}: {decoded_output}")
|
|
|
+ #passed to popup through app.py and label_dailog.py
|
|
|
+ self.decoded_barcode = decoded_output
|
|
|
+ # (Optional) Use the decoded output for updating UI or further processing
|
|
|
+ # Example: Print or log the decoded text
|
|
|
+ # print(f"Decoded output: {decoded_output}")
|
|
|
+ except Exception as e:
|
|
|
+ logger.error(f"Error during decoding for bbox points {bbox_points}: {e}")
|
|
|
+
|
|
|
+ ## pass each shape points to decode function here, do the decodign and give back the output here and print the output
|
|
|
if self.pred_bbox_points:
|
|
|
self.draw_pred = True
|
|
|
self.finalise()
|
|
@@ -1212,7 +1319,6 @@ class Canvas(QtWidgets.QWidget):
|
|
|
if all_segmentation_results:
|
|
|
logger.info(f"Segmentation results for all detections: {all_segmentation_results}")
|
|
|
|
|
|
-
|
|
|
def keyReleaseEvent(self, ev):
|
|
|
modifiers = ev.modifiers()
|
|
|
if self.drawing():
|
|
@@ -1264,10 +1370,6 @@ class Canvas(QtWidgets.QWidget):
|
|
|
|
|
|
def loadPixmap(self, pixmap, clear_shapes=True):
|
|
|
self.pixmap = pixmap
|
|
|
- if self._detection_model:
|
|
|
- self._detection_model.set_image(
|
|
|
- image=labelme.utils.img_qt_to_arr(self.pixmap.toImage())
|
|
|
- )
|
|
|
if clear_shapes:
|
|
|
self.shapes = []
|
|
|
self.update()
|