save into file success after label a box and key

PPOCRLabel/PPOCRLabel.py

@@ -1106,24 +1106,25 @@ def format_shape(s):
  line_color=s.line_color.getRgb(),
  fill_color=s.fill_color.getRgb(),
  points=[(int(p.x()), int(p.y())) for p in s.points], # QPonitF
- # add chris
- difficult=s.difficult) # bool
+ difficult=s.difficult,
+ key_cls=s.key_cls) # bool
 
- shapes = [] if mode == 'Auto' else \
- [format_shape(shape) for shape in self.canvas.shapes if shape.line_color != DEFAULT_LOCK_COLOR]
+ if mode == 'Auto':
+ shapes = []
+ else:
+ shapes = [format_shape(shape) for shape in self.canvas.shapes if shape.line_color != DEFAULT_LOCK_COLOR]
  # Can add differrent annotation formats here
  for box in self.result_dic:
- trans_dic = {"label": box[1][0], "points": box[0], 'difficult': False}
+ trans_dic = {"label": box[1][0], "points": box[0], "difficult": False, "key_cls": "None"}
  if trans_dic["label"] == "" and mode == 'Auto':
  continue
  shapes.append(trans_dic)
 
  try:
  trans_dic = []
  for box in shapes:
- trans_dic.append(
- {"transcription": box['label'], "points": box['points'],
- "difficult": box['difficult'], "key": "None"})
+ trans_dic.append({"transcription": box['label'], "points": box['points'],
+ "difficult": box['difficult'], "key_cls": box['key_cls']})
  self.PPlabel[annotationFilePath] = trans_dic
  if mode == 'Auto':
  self.Cachelabel[annotationFilePath] = trans_dic
@@ -1205,8 +1206,21 @@ def newShape(self, value=True):
 
  if text is not None:
  self.prevLabelText = self.stringBundle.getString('tempLabel')
- # generate_color = generateColorByText(text)
- shape = self.canvas.setLastLabel(text, None, None) # generate_color, generate_color
+
+ if self.kie_mode:
+ key_text, _ = self.keyDialog.popUp(self.key_previous_text)
+ if key_text is not None:
+ shape = self.canvas.setLastLabel(text, None, None, key_text) # generate_color, generate_color
+ self.key_previous_text = key_text
+ if not self.keyList.findItemsByLabel(key_text):
+ item = self.keyList.createItemFromLabel(key_text)
+ self.keyList.addItem(item)
+ rgb = self._get_rgb_by_label(key_text, self.kie_mode)
+ self.keyList.setItemLabel(item, key_text, rgb)
+ self.keyDialog.addLabelHistory(key_text)
+ else:
+ shape = self.canvas.setLastLabel(text, None, None, None) # generate_color, generate_color
+
  self.addLabel(shape)
  if self.beginner(): # Switch to edit mode.
  self.canvas.setEditing(True)
@@ -1221,17 +1235,6 @@ def newShape(self, value=True):
  # self.canvas.undoLastLine()
  self.canvas.resetAllLines()
 
- if self.kie_mode:
- key_text, flags = self.keyDialog.popUp(self.key_previous_text)
- if key_text is not None:
- self.key_previous_text = key_text
- if not self.keyList.findItemsByLabel(key_text):
- item = self.keyList.createItemFromLabel(key_text)
- self.keyList.addItem(item)
- rgb = self._get_rgb_by_label(key_text, self.kie_mode)
- self.keyList.setItemLabel(item, key_text, rgb)
- self.keyDialog.addLabelHistory(key_text)
-
  def _update_shape_color(self, shape):
  r, g, b = self._get_rgb_by_label(shape.label)
  shape.line_color = QtGui.QColor(r, g, b)
@@ -1438,13 +1441,13 @@ def showBoundingBoxFromPPlabel(self, filePath):
  for box in self.canvas.lockedShapes:
  if self.canvas.isInTheSameImage:
  shapes.append((box['transcription'], [[s[0] * width, s[1] * height] for s in box['ratio']],
- DEFAULT_LOCK_COLOR, box['key'], box['difficult']))
+ DEFAULT_LOCK_COLOR, box['key_cls'], box['difficult']))
  else:
  shapes.append(('锁定框：待检测', [[s[0] * width, s[1] * height] for s in box['ratio']],
- DEFAULT_LOCK_COLOR, box['key'], box['difficult']))
+ DEFAULT_LOCK_COLOR, box['key_cls'], box['difficult']))
  if imgidx in self.PPlabel.keys():
  for box in self.PPlabel[imgidx]:
- shapes.append((box['transcription'], box['points'], None, box['key'], box['difficult']))
+ shapes.append((box['transcription'], box['points'], None, box['key_cls'], box['difficult']))
 
  self.loadLabels(shapes)
  self.canvas.verified = False
@@ -2135,7 +2138,8 @@ def saveRecResult(self):
  try:
  img = cv2.imread(key)
  for i, label in enumerate(self.PPlabel[idx]):
- if label['difficult']: continue
+ if label['difficult']:
+ continue
  img_crop = get_rotate_crop_image(img, np.array(label['points'], np.float32))
  img_name = os.path.splitext(os.path.basename(idx))[0] + '_crop_' + str(i) + '.jpg'
  cv2.imwrite(crop_img_dir + img_name, img_crop)
@@ -2202,8 +2206,9 @@ def format_shape(s):
  line_color=s.line_color.getRgb(),
  fill_color=s.fill_color.getRgb(),
  ratio=[[int(p.x()) / width, int(p.y()) / height] for p in s.points], # QPonitF
- # add chris
- difficult=s.difficult) # bool
+ difficult=s.difficult,# bool
+ key_cls=s.key_cls,# bool
+ )
 
  # lock
  if len(self.canvas.lockedShapes) == 0:
@@ -2214,7 +2219,8 @@ def format_shape(s):
  trans_dic = []
  for box in shapes:
  trans_dic.append({"transcription": box['label'], "ratio": box['ratio'],
- "difficult": box['difficult'], "key": "None" if "key" not in box else box["key"]})
+ "difficult": box['difficult'],
+ "key_cls": "None" if "key_cls" not in box else box["key_cls"]})
  self.canvas.lockedShapes = trans_dic
  self.actions.save.setEnabled(True)
 

PPOCRLabel/libs/canvas.py

@@ -783,14 +783,18 @@ 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, key_cls=None):
  assert text
  self.shapes[-1].label = text
  if line_color:
  self.shapes[-1].line_color = line_color
 
  if fill_color:
  self.shapes[-1].fill_color = fill_color
+
+ if key_cls:
+ self.shapes[-1].key_cls = key_cls
+
  self.storeShapes()
 
  return self.shapes[-1]

PPOCRLabel/libs/shape.py

@@ -46,12 +46,13 @@ class Shape(object):
  point_size = 8
  scale = 1.0
 
- def __init__(self, label=None, line_color=None, difficult=False, paintLabel=False):
+ def __init__(self, label=None, line_color=None, difficult=False, key_cls="None", paintLabel=False):
  self.label = label
  self.points = []
  self.fill = False
  self.selected = False
  self.difficult = difficult
+ self.key_cls = key_cls
  self.paintLabel = paintLabel
  self.locked = False
  self.direction = 0
@@ -224,6 +225,7 @@ def copy(self):
  if self.fill_color != Shape.fill_color:
  shape.fill_color = self.fill_color
  shape.difficult = self.difficult
+ shape.key_cls = self.key_cls
  return shape
 
  def __len__(self):