Merge pull request #39 from Wei-JL/dev

增加多点标注且适应多点框模型重识别

PPOCRLabel.py

@@ -1866,6 +1866,37 @@ def autoRecognition(self):
         self.setDirty()
         self.saveCacheLabel()
 
+    def gen_quad_from_poly(self, poly):
+        """
+        Generate min area quad from poly.
+        """
+        point_num = poly.shape[0]
+        min_area_quad = np.zeros ((4, 2), dtype=np.float32)
+        rect = cv2.minAreaRect (poly.astype (
+            np.int32))  # (center (x,y), (width, height), angle of rotation)
+        box = np.array (cv2.boxPoints (rect))
+
+        first_point_idx = 0
+        min_dist = 1e4
+        for i in range (4):
+            dist = np.linalg.norm (box[(i + 0) % 4] - poly[0]) + \
+                   np.linalg.norm (box[(i + 1) % 4] - poly[point_num // 2 - 1]) + \
+                   np.linalg.norm (box[(i + 2) % 4] - poly[point_num // 2]) + \
+                   np.linalg.norm (box[(i + 3) % 4] - poly[-1])
+            if dist < min_dist:
+                min_dist = dist
+                first_point_idx = i
+        for i in range(4):
+            min_area_quad[i] = box[(first_point_idx + i) % 4]
+
+        bbox_new = min_area_quad.tolist()
+        bbox = []
+
+        for box in bbox_new:
+            box = list(map(int, box))
+            bbox.append(box)
+
+        return bbox
 
     def reRecognition(self):
         img = cv2.imread(self.filePath)
@@ -1875,8 +1906,17 @@ def reRecognition(self):
             rec_flag = 0
             for shape in self.canvas.shapes:
                 box = [[int(p.x()), int(p.y())] for p in shape.points]
-                assert len(box) == 4
+                temp = 0
+                if len(box) > 4:
+                    temp = box
+                    box = np.array(box)
+                    box = self.gen_quad_from_poly(box)
+
+                assert len(box) >= 4
                 img_crop = get_rotate_crop_image(img, np.array(box, np.float32))
+                if not temp == 0:
+                    box = temp
+
                 if img_crop is None:
                     msg = 'Can not recognise the detection box in ' + self.filePath + '. Please change manually'
                     QMessageBox.information(self, "Information", msg)
@@ -1910,8 +1950,17 @@ def singleRerecognition(self):
         img = cv2.imread(self.filePath)
         for shape in self.canvas.selectedShapes:
             box = [[int(p.x()), int(p.y())] for p in shape.points]
-            assert len(box) == 4
+            temp = 0
+            if len(box) > 4:
+                temp = box
+                box = np.array(box)
+                box = self.gen_quad_from_poly(box)
+
+            assert len(box) >= 4
             img_crop = get_rotate_crop_image(img, np.array(box, np.float32))
+            if not temp == 0:
+                box = temp
+
             if img_crop is None:
                 msg = 'Can not recognise the detection box in ' + self.filePath + '. Please change manually'
                 QMessageBox.information(self, "Information", msg)

libs/canvas.py

@@ -233,7 +233,7 @@ def mouseMoveEvent(self, ev):
                 self.hVertex, self.hShape = index, shape
                 shape.highlightVertex(index, shape.MOVE_VERTEX)
                 self.overrideCursor(CURSOR_POINT)
-                self.setToolTip("Click & drag to move point")
+                self.setToolTip("Click & drag to move point") #move point
                 self.setStatusTip(self.toolTip())
                 self.update()
                 break
@@ -272,6 +272,14 @@ def mousePressEvent(self, ev):
                         assert len(self.current.points) == 1
                         self.current.points = self.line.points
                         self.finalise()
+
+                    if self.canCloseShape() and len(self.current) > 3 and self.current[0].x() + 2 >= pos.x() >= \
+                            self.current[0].x() - 2 and self.current[0].y() + 2 >= pos.y() >= self.current[0].y() - 2:
+                        print('鼠标单击事件')
+                        if len(self.current) > 4:
+                            self.current.popPoint() # Eliminate the extra point from the last click.
+                        self.finalise()
+
                 elif not self.outOfPixmap(pos):
                     # Create new shape.
                     self.current = Shape()
@@ -488,6 +496,7 @@ def boundedMoveVertex(self, pos):
             shape.moveVertexBy(lindex, lshift)
 
         else:
+            #move point
             shape.moveVertexBy(index, shiftPos)
 
 
@@ -650,7 +659,7 @@ def outOfPixmap(self, p):
 
     def finalise(self):
         assert self.current
-        if self.current.points[0] == self.current.points[-1]:
+        if len(self.current) < 4 and self.current.points[0] == self.current.points[-1]:
             # print('finalse')
             self.current = None
             self.drawingPolygon.emit(False)
@@ -764,7 +773,7 @@ def moveOutOfBound(self, step):
         points = [p1+p2 for p1, p2 in zip(self.selectedShape.points, [step]*4)]
         return True in map(self.outOfPixmap, points)
 
-    def setLastLabel(self, text, line_color  = None, fill_color = None):
+    def setLastLabel(self, text, line_color=None, fill_color=None):
         assert text
         self.shapes[-1].label = text
         if line_color:

libs/shape.py

@@ -77,7 +77,7 @@ def close(self):
         self._closed = True
 
     def reachMaxPoints(self):
-        if len(self.points) >= 4:
+        if len(self.points) >= 65535:
             return True
         return False