labelme/__main__.py

@@ -109,6 +109,11 @@ def main():
         help="model weight for barcode detection",
         default=argparse.SUPPRESS,
     )
+    parser.add_argument(
+        "--segmentation_model",
+        type=str,
+        help="Path to the segmentation model"
+    )
     args = parser.parse_args()
 
     if args.version:
@@ -146,6 +151,7 @@ def main():
     config_file_or_yaml = config_from_args.pop("config")
     config = get_config(config_file_or_yaml, config_from_args)
     config["model"]=config_from_args.pop("model")
+    config["segmentation_model"]=config_from_args.pop("segmentation_model")
     if not config["labels"] and config["validate_label"]:
         logger.error(
             "--labels must be specified with --validatelabel or "

labelme/ai/__init__.py

@@ -2,12 +2,15 @@ import gdown
 from .efficient_sam import EfficientSam
 from .segment_anything_model import SegmentAnythingModel
 from .barcode_model import BarcodePredictModel
+
 class BarcodePredict(BarcodePredictModel):
     name="BarcodePredict(ov)"
-    def __init__(self,model_path=None):
+    def __init__(self, detection_model_path=None, segmentation_model_path=None):
         super().__init__(
-            model_path=model_path
+            detection_model_path=detection_model_path,
+            segmentation_model_path=segmentation_model_path
         )
+
 class SegmentAnythingModelVitB(SegmentAnythingModel):
     name = "SegmentAnything (speed)"
 

labelme/ai/barcode_model.py

@@ -1,81 +0,0 @@
-import collections
-import threading
-import numpy as np
-import openvino as ov
-import os.path as osp
-import cv2
-from ..logger import logger
-from . import _utils
-import copy
-class BarcodePredictModel:
-    def __init__(self,model_path):
-        self.ie =  ov.Core()
-        bin_path = osp.splitext(model_path)[0] + '.bin'
-        self.net = self.ie.read_model(model_path, bin_path)
-        self.device = 'cpu'
-        self._lock = threading.Lock()
-        self._image_embedding_cache = collections.OrderedDict()
-        self.ppp=ov.preprocess.PrePostProcessor(self.net)
-        self.ppp.input().tensor().set_shape([1,480,480,1]).set_element_type(ov.Type.u8).set_layout(ov.Layout('NHWC'))
-        self.ppp.input().model().set_layout(ov.Layout('NCHW'))
-        # self.ppp.output().tensor().set_element_type(ov.Type.f32)
-        self.net=self.ppp.build()
-        self.sess = self.ie.compile_model(model=self.net, device_name=self.device.upper())
-        self.predict_session=self.sess.create_infer_request()
-
-    def set_image(self,image:np.ndarray):
-        with self._lock:
-            #keep ratio is false
-            self.raw_width=image.shape[1]
-            self.raw_height=image.shape[0]
-            resized_image=cv2.resize(image,(480,480))
-            input_tensor=np.expand_dims(resized_image,2)
-            input_tensor=np.expand_dims(input_tensor,0)
-            self._image=input_tensor
-            self._thread = threading.Thread(
-                target=self._compute_and_cache_image_embedding
-            )
-            self._thread.start()
-    def _compute_and_cache_image_embedding(self):
-        with self._lock:
-            logger.debug("Computing image embedding using BarcodePredict...")
-            #we dont need to transfer because we use mono camera
-            self._result =  self.predict_session.infer(self._image)
-            logger.debug("Done computing image embedding.")
-    def _get_image_embedding(self):
-        if self._thread is not None:
-            self._thread.join()
-            self._thread=None
-        with self._lock:
-            new_result=self._result
-            return new_result
-    def predict_mask_from_points(self,points=None,point_labels=None):
-        return _collect_result_from_output(
-            outputs=self._get_image_embedding(),
-            raw_width=self.raw_width,
-            raw_height=self.raw_height,
-        )
-    def predict_polygon_from_points(self,points=None,point_labels=None):
-        result_list=self.predict_mask_from_points(points,point_labels)
-        return result_list
-def _collect_result_from_output(outputs,raw_width,raw_height):
-    outputs=np.squeeze(outputs[0],axis=0)
-    point_list=[]
-    thresh_hold=0.5
-
-    for bbox_info in outputs:
-        if(bbox_info[4]>thresh_hold):
-            ratio_x=raw_width/480
-            ratio_y=raw_height/480
-            #0->x1 1->y1 2->x2 3->y2
-            x1=max(min(int(ratio_x*bbox_info[0]),raw_width-1),0)
-            y1=max(min(int(ratio_y*bbox_info[1]),raw_height-1),0)
-            x2=max(min(int(ratio_x*bbox_info[2]),raw_width-1),0)
-            y2=max(min(int(ratio_y*bbox_info[3]),raw_height-1),0)
-            point_xy=[]
-            point_xy.append([x1,y1])
-            point_xy.append([x2,y1])
-            point_xy.append([x2,y2])
-            point_xy.append([x1,y2])
-            point_list.append(point_xy)
-    return point_list

labelme/app.py

@@ -166,7 +166,8 @@ class MainWindow(QtWidgets.QMainWindow):
             double_click=self._config["canvas"]["double_click"],
             num_backups=self._config["canvas"]["num_backups"],
             crosshair=self._config["canvas"]["crosshair"],
-            model_path=self._config["model"],
+            detection_model_path=self._config["model"],
+            segmentation_model_path=self._config["segmentation_model"],
         )
         self.canvas.zoomRequest.connect(self.zoomRequest)
         self.canvas.mouseMoved.connect(
@@ -1240,8 +1241,6 @@ class MainWindow(QtWidgets.QMainWindow):
         self.actions.edit.setEnabled(n_selected)
 
     def addLabel(self, shape):
-        if(shape.label is None):
-            return
         if shape.group_id is None:
             text = shape.label
         else:

labelme/config/__init__.py

@@ -69,7 +69,7 @@ def get_config(config_file_or_yaml=None, config_from_args=None):
         update_dict(config, config_from_yaml, validate_item=validate_config_item)
 
     # 3. command line argument or specified config file
-    print(config_from_args)
+    # print(config_from_args)
     if config_from_args is not None:
         update_dict(config, config_from_args, validate_item=validate_config_item)
 

labelme/widgets/canvas.py

@@ -8,7 +8,8 @@ import labelme.utils
 from labelme import QT5
 from labelme.logger import logger
 from labelme.shape import Shape
-
+import numpy as np
+import cv2
 # TODO(unknown):
 # - [maybe] Find optimal epsilon value.
 
@@ -41,7 +42,8 @@ class Canvas(QtWidgets.QWidget):
 
     def __init__(self, *args, **kwargs):
         self.epsilon = kwargs.pop("epsilon", 10.0)
-        self.model_path=kwargs.pop("model_path",None)
+        self.detection_model_path=kwargs.pop("detection_model_path",None)
+        self.segmentation_model_path=kwargs.pop("segmentation_model_path",None)
         self.double_click = kwargs.pop("double_click", "close")
         if self.double_click not in [None, "close"]:
             raise ValueError(
@@ -106,7 +108,8 @@ class Canvas(QtWidgets.QWidget):
         self.setFocusPolicy(QtCore.Qt.WheelFocus)
 
         self._ai_model = None
-
+        self._detection_model = None
+        self._segmentation_model = None
     def fillDrawing(self):
         return self._fill_drawing
 
@@ -132,16 +135,36 @@ class Canvas(QtWidgets.QWidget):
             raise ValueError("Unsupported createMode: %s" % value)
         self._createMode = value
 
-    def initializeAiModel(self,name,weight_path=None):
-        if name not in [model.name for model in labelme.ai.MODELS]:
-            raise ValueError("Unsupported ai model: %s" % name)
-        model = [model for model in labelme.ai.MODELS if model.name == name][0]
+    def initializeBarcodeModel(self, detection_model_path, segmentation_model_path=None):
+        if not detection_model_path:
+            raise ValueError("Detection model path is required.")
 
-        if self._ai_model is not None and self._ai_model.name == model.name:
-            logger.debug("AI model is already initialized: %r" % model.name)
-        else:
-            logger.debug("Initializing AI model: %r" % model.name)
-            self._ai_model = model(weight_path)
+        logger.debug("Initializing only detection model: %r" % "BarcodePredictModel")
+        self._detection_model = labelme.ai.BarcodePredictModel(detection_model_path)
+
+        if segmentation_model_path:
+            logger.debug("Initializing barcode detection  & Segmentation model: %r" % "BarcodePredictModel")
+            self._segmentation_model = labelme.ai.BarcodePredictModel(
+                detection_model_path, segmentation_model_path
+            )
+        if self.pixmap is None:
+            logger.warning("Pixmap is not set yet")
+            return
+
+        self._detection_model.set_image(
+            image=labelme.utils.img_qt_to_arr(self.pixmap.toImage())
+        )
+
+    def initializeAiModel(self, name, weight_path=None):
+        if self._ai_model is not None:
+            logger.debug("AI model is already initialized.")
+            return
+
+        if name not in [model.name for model in labelme.ai.MODELS]:
+            raise ValueError("Unsupported AI model: %s" % name)
+        model_class = [model for model in labelme.ai.MODELS if model.name == name][0]
+        logger.debug(f"Initializing AI model: {name}")
+        self._ai_model = model_class(weight_path)
 
         if self.pixmap is None:
             logger.warning("Pixmap is not set yet")
@@ -768,7 +791,7 @@ class Canvas(QtWidgets.QWidget):
                 point=self.line.points[1],
                 label=self.line.point_labels[1],
             )
-            points = self._ai_model.predict_polygon_from_points(
+            points = self._detection_model.predict_polygon_from_points(
                 points=[[point.x(), point.y()] for point in drawing_shape.points],
                 point_labels=drawing_shape.point_labels,
             )
@@ -787,7 +810,7 @@ class Canvas(QtWidgets.QWidget):
                 point=self.line.points[1],
                 label=self.line.point_labels[1],
             )
-            mask = self._ai_model.predict_mask_from_points(
+            mask = self._detection_model.predict_mask_from_points(
                 points=[[point.x(), point.y()] for point in drawing_shape.points],
                 point_labels=drawing_shape.point_labels,
             )
@@ -826,7 +849,7 @@ class Canvas(QtWidgets.QWidget):
         if self.createMode == "ai_polygon":
             # convert points to polygon by an AI model
             assert self.current.shape_type == "points"
-            points = self._ai_model.predict_polygon_from_points(
+            points = self._detection_model.predict_polygon_from_points(
                 points=[[point.x(), point.y()] for point in self.current.points],
                 point_labels=self.current.point_labels,
             )
@@ -838,7 +861,7 @@ class Canvas(QtWidgets.QWidget):
         elif self.createMode == "ai_mask":
             # convert points to mask by an AI model
             assert self.current.shape_type == "points"
-            mask = self._ai_model.predict_mask_from_points(
+            mask = self._detection_model.predict_mask_from_points(
                 points=[[point.x(), point.y()] for point in self.current.points],
                 point_labels=self.current.point_labels,
             )
@@ -1000,18 +1023,28 @@ class Canvas(QtWidgets.QWidget):
             elif modifiers == QtCore.Qt.AltModifier:
                 self.snapping = False
             elif key == QtCore.Qt.Key_V:
-                if self._ai_model is None and self.model_path is not None:
-                    print(f"model got initialized")
-                    self.initializeAiModel("BarcodePredict(ov)",self.model_path)
+                if self._detection_model is None:
+                    logger.info(f"Initializing AI model")
+                    self.initializeBarcodeModel(self.detection_model_path, self.segmentation_model_path)
+ 
                 self.current = Shape(
-                        shape_type="points"
-                        if self.createMode in ["ai_polygon", "ai_mask"]
-                        else self.createMode
-                    )
-                if self._ai_model is not None:
-                    self.pred_bbox_points=self._ai_model.predict_polygon_from_points()
-                    self.draw_pred=True
-                    self.finalise()         
+                    shape_type="points" if self.createMode in ["ai_polygon", "ai_mask"] else self.createMode
+                )
+                        
+                if self._detection_model:
+                    if self._segmentation_model is None:
+                        logger.info(f"Performing detection only.")
+                        # Get prediction from model
+                        self.pred_bbox_points = self._detection_model.predict_polygon_from_points()
+                        print("Predicted Bounding Box Points:", self.pred_bbox_points)                        
+                        if self.pred_bbox_points:
+                            self.draw_pred = True
+                            self.finalise()
+                        else:
+                            print("No bounding boxes detected.")    
+                    else:
+                            logger.info(f"Performing detection and segmentation.")
+                            self.detect_and_segment() 
         elif self.editing():
             if key == QtCore.Qt.Key_Up:
                 self.moveByKeyboard(QtCore.QPointF(0.0, -MOVE_SPEED))
@@ -1022,8 +1055,163 @@ 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):
+        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]
+    
+    def detect_and_segment(self):
+        """
+        Perform detection and segmentation (if both models are available).
+        """
+        logger.info("Performing detection and segmentation.")
+        self.current = Shape(
+            shape_type="points" if self.createMode in ["ai_polygon", "ai_mask"] else self.createMode
+        )
+
+        # Step 1: detection bounding box points
+        detection_results = self._detection_model.predict_polygon_from_points()
+
+        if not detection_results or len(detection_results) == 0:
+            logger.warning("No detection found")
+            return
+
+        logger.debug(f"Detection results: {detection_results}")
+        all_segmentation_results = []
+        rotated = False
+        # Step 2: Loop through each detection result since there are multiple per image
+        for detection_idx, detection_result in enumerate(detection_results):
+            logger.debug(f"Processing detection {detection_idx + 1}/{len(detection_results)}")
+
+            try:
+                #extracting
+                x_coords = [point[0] for point in detection_result]
+                y_coords = [point[1] for point in detection_result]
+                
+                #min and max values for x and y
+                x_min, x_max = min(x_coords), max(x_coords)
+                y_min, y_max = min(y_coords), max(y_coords)
+                
+                logger.debug(f"Bounding box for detection {detection_idx + 1} - x_min: {x_min}, y_min: {y_min}, x_max: {x_max}, y_max: {y_max}")
+            except Exception as e:
+                logger.error(f"Error extracting bounding box coordinates for detection {detection_idx + 1}: {e}")
+                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)
+
+            # Step 3: Cropping image based on detection output
+            try:
+                
+                cropped_image = self.pixmap.toImage().copy(x_min, y_min, x_max - x_min, y_max - y_min)
+                cropped_image = labelme.utils.img_qt_to_arr(cropped_image)
+
+                orig_height, orig_width = cropped_image.shape[:2]
+                logger.debug(f"Original height: {orig_height}, Original width: {orig_width}")
+                # if the height is greater than the width we rotate for segmentaion
+                if orig_height > orig_width:  
+                    cropped_image = cv2.rotate(cropped_image, cv2.ROTATE_90_CLOCKWISE)
+                    logger.debug(f"Rotated cropped image for detection {detection_idx + 1} due to height > width.")
+                    orig_cropped_shape = cropped_image.shape[:2]
+                    rotated = True
+
+                else: 
+                    rotated = False
+                # Save crop image
+                cv2.imwrite(f"cropped_image_{detection_idx + 1}.png", cropped_image)
+                logger.debug(f"Saved cropped image for detection {detection_idx + 1}: {cropped_image.shape}")
+
+                # logger.debug(f"Cropped image shape for detection {detection_idx + 1}: {cropped_image.shape}")
+            except Exception as e:
+                logger.error(f"Error cropping the image for detection {detection_idx + 1}: {e}")
+                continue
+
+            # Step 4: Resize the cropped image to match segmentation input size (1 64 256)
+            try:
+                orig_cropped_shape = cropped_image.shape[:2]  # Save the original cropped image size
+                preprocessed_img = self._detection_model.preprocess_image(cropped_image, for_segmentation=True)
+                logger.debug(f"Preprocessed image shape for segmentation detection {detection_idx + 1}: {preprocessed_img.shape}")
+            except Exception as e:
+                logger.error(f"Error preprocessing the image for segmentation for detection {detection_idx + 1}: {e}")
+                continue
+
+            # Step 5: inference on segmentation model on cropped image
+            try:
+                seg_result = self._segmentation_model.segmentation_sess.infer_new_request({'x': preprocessed_img})
+                logger.debug(f"Segmentation model inference completed for detection {detection_idx + 1}.")
+            except Exception as e:
+                logger.error(f"Error during segmentation model inference for detection {detection_idx + 1}: {e}")
+                continue
 
+            # Step 6: Convert binary mask to polygon (contours)
+            try:
+                mask = seg_result['save_infer_model/scale_0.tmp_0']  #model output name
+                mask = mask.squeeze()  # Remove batch dimension, should result in (64, 256)
+                logger.debug(f"Segmentation mask shape for detection {detection_idx + 1}: {mask.shape}")
+
+                # 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)
+                if rotated:
+                    cropped_image = cv2.rotate(cropped_image, cv2.ROTATE_90_COUNTERCLOCKWISE)
+                    mask = cv2.rotate(mask, cv2.ROTATE_90_COUNTERCLOCKWISE)
+                    rotated_cropped_shape = cropped_image.shape[:2]
+                
+                # cv2.imwrite(f"segmentation_mask_{detection_idx + 1}.png", mask)
+                logger.debug(f"Saved segmentation mask for detection {detection_idx + 1}.")
+                
+                # Step 7: Find contours
+                contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                logger.debug(f"Found {len(contours)} contours in the mask for detection {detection_idx + 1}.")
+
+                if len(contours) > 0:
+                    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
+                        ]
+                    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}")
+                    if self.pred_bbox_points:
+                        self.draw_pred = True
+                        self.finalise()
+                    else:
+                        logger.info(f"No bounding boxes detected for detection {detection_idx + 1}.")
+                        
+                    # Collect the segmentation result
+                    all_segmentation_results.append(self.pred_bbox_points)
+            except Exception as e:
+                logger.error(f"Error creating the polygon shape for detection {detection_idx + 1}: {e}")
+
+            # **Reset critical variables after each detection**:
+            self.pred_bbox_points = None
+            self.draw_pred = False
+
+        # You now have a list of segmentation results for all detections
+        if all_segmentation_results:
+            logger.info(f"Segmentation results for all detections: {all_segmentation_results}")
+
 
     def keyReleaseEvent(self, ev):
         modifiers = ev.modifiers()
@@ -1076,8 +1264,8 @@ class Canvas(QtWidgets.QWidget):
 
     def loadPixmap(self, pixmap, clear_shapes=True):
         self.pixmap = pixmap
-        if self._ai_model:
-            self._ai_model.set_image(
+        if self._detection_model:
+            self._detection_model.set_image(
                 image=labelme.utils.img_qt_to_arr(self.pixmap.toImage())
             )
         if clear_shapes: