python实现简单图片物体标注工具

本文实例为大家分享了python实现简单图片物体标注工具的具体代码，供大家参考，具体内容如下

# coding: utf-8

"""
物体检测标注小工具
基本思路：
对要标注的图像建立一个窗口循环，然后每次循环的时候对图像进行一次复制，
鼠标在画面上画框的操作、画好的框的相关信息在全局变量中保存，
并且在每个循环中根据这些信息，在复制的图像上重新画一遍，然后显示这份复制的图像。
简化的设计过程：
1、输入是一个文件夹的路径，包含了所需标注物体框的图片。
如果图片中标注了物体，则生成一个相同名称加额外后缀_bbox的文件，来保存标注信息。
2、标注的方式：按下鼠标左键选择物体框的左上角，松开鼠标左键选择物体框的右下角，
按下鼠标右键删除上一个标注好的物体框。
所有待标注物体的类别和标注框颜色由用户自定义。
如果没有定义则默认只标注一种物体，定义该物体名称为Object。
3、方向键 ← 和 → 键用来遍历图片， ↑ 和 ↓ 键用来选择当前要标注的物体，
Delete键删除一种脏图片和对应的标注信息。
自定义标注物体和颜色的信息用一个元组表示
第一个元素表示物体名字
第二个元素表示BGR颜色的tuple或者代表标注框坐标的元祖
利用repr()保存和eval()读取
"""

"""
一些说明：
1. 标注相关的物体标签文件即 .labels 结尾的文件，需要与所选文件夹添加到同一个根目录下
一定要注意这一点，否则无法更新标注物体的类型标签，致使从始至终都只有一个默认物体出现
我就是这个原因，拖了两三天才整好，当然也顺便仔细的读了这篇代码。同时也学习了@staticmethod以及相应Python的decorator的知识。
可以说，在曲折中前进才是棒的。
2. .labels文件为预设物体标签文件，其内容具体格式为:
'object1', (B, G, R)
'object2', (B, G, R)
'object3', (B, G, R)……
具体见文后图片。
3. 最后生成的标注文件，在文后会有，到时再进行解释。
"""

import os
import cv2
# tkinter是Python内置的简单GUI库，实现打开文件夹、确认删除等操作十分方便
from tkMessageBox import askyesno
# 定义标注窗口的默认名称
WINDOW_NAME = 'Simple Bounding Box Labeling Tool'
# 定义画面刷新帧率
FPS = 24
# 定义支持的图像格式
SUPPORTED_FORMATS = ['jpg', 'jpeg', 'png']
# 定义默认物体框的名字为Object，颜色为蓝色，当没有用户自定义物体时，使用该物体
DEFAULT_COLOR = {'Object': (255, 0, 0)}
# 定义灰色，用于信息显示的背景和未定义物体框的显示
COLOR_GRAY = (192, 192, 192)
# 在图像下方多处BAR_HEIGHT的区域，用于显示信息
BAR_HEIGHT = 16
# 上下左右,DELETE键对应的cv2.waitKey()函数的返回值
KEY_UP = 2490368
KEY_DOWN = 2621440
KEY_LEFT = 2424832
KEY_RIGHT = 2555904
KEY_DELETE = 3014656
# 空键用于默认循环
KEY_EMPTY = 0
get_bbox_name = '{}.bbox'.format

# 定义物体框标注工具类
class SimpleBBoxLabeling:
def __init__(self, data_dir, fps=FPS, windown_name=WINDOW_NAME):
 self._data_dir = data_dir
 self.fps = fps
 self.window_name = windown_name if windown_name else WINDOW_NAME

# pt0 是正在画的左上角坐标, pt1 是鼠标所在坐标
 self._pt0 = None
 self._pt1 = None
 # 表明当前是否正在画框的状态标记
 self._drawing = False
 # 当前标注物体的名称
 self._cur_label = None
 # 当前图像对应的所有已标注框
 self._bboxes = []
 # 如果有用户自己定义的标注信息则读取，否则使用默认的物体和颜色
 label_path = '{}.labels'.format(self._data_dir)
 self.label_colors = DEFAULT_COLOR if not os.path.exists(label_path) else self.load_labels(label_path)
 # self.label_colors = self.load_labels(label_path)
 # 获取已经标注的文件列表和未标注的文件列表
 imagefiles = [x for x in os.listdir(self._data_dir) if x[x.rfind('.') + 1:].lower() in SUPPORTED_FORMATS]
 labeled = [x for x in imagefiles if os.path.exists(get_bbox_name(x))]
 to_be_labeled = [x for x in imagefiles if x not in labeled]

# 每次打开一个文件夹，都自动从还未标注的第一张开始
 self._filelist = labeled + to_be_labeled
 self._index = len(labeled)
 if self._index > len(self._filelist) - 1:
  self._index = len(self._filelist) - 1

# 鼠标回调函数
def _mouse_ops(self, event, x, y, flags, param):
 # 按下左键，坐标为左上角，同时表示开始画框，改变drawing，标记为True
 if event == cv2.EVENT_LBUTTONDOWN:
  self._drawing = True
  self._pt0 = (x, y)
 # 松开左键，表明画框结束，坐标为有效较并保存，同时改变drawing，标记为False
 elif event == cv2.EVENT_LBUTTONUP:
  self._drawing = False
  self._pt1 = (x, y)
  self._bboxes.append((self._cur_label, (self._pt0, self._pt1)))
 # 实时更新右下角坐标
 elif event == cv2.EVENT_MOUSEMOVE:
  self._pt1 = (x, y)
 # 按下鼠标右键删除最近画好的框
 elif event == cv2.EVENT_RBUTTONUP:
  if self._bboxes:
   self._bboxes.pop()

# 清除所有标注框和当前状态
def _clean_bbox(self):
 self._pt0 = None
 self._pt1 = None
 self._drawing = False
 self._bboxes = []

# 画标注框和当前信息的函数
def _draw_bbox(self, img):
 # 在图像下方多出BAR_HEIGHT的区域，显示物体信息
 h, w = img.shape[:2]
 canvas = cv2.copyMakeBorder(img, 0, BAR_HEIGHT, 0, 0, cv2.BORDER_CONSTANT, value=COLOR_GRAY)
 # 正在标注的物体信息，如果鼠标左键已经按下，则像是两个点坐标，否则显示当前待标注物体的名
 label_msg = '{}: {}, {}'.format(self._cur_label, self._pt0, self._pt1) \
  if self._drawing \
  else 'Current label: {}'.format(self._cur_label)
 # 显示当前文件名，文件个数信息
 msg = '{}/{}: {} | {}'.format(self._index + 1, len(self._filelist), self._filelist[self._index], label_msg)
 cv2.putText(canvas, msg, (1, h+12), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
 # 画出已经标好的框和对应名字
 for label, (bpt0, bpt1) in self._bboxes:
  label_color = self.label_colors[label] if label in self.label_colors else COLOR_GRAY
  cv2.rectangle(canvas, bpt0, bpt1, label_color, thickness=2)
  cv2.putText(canvas, label, (bpt0[0]+3, bpt0[1]+15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, label_color, 2)
 # 画正在标注的框和对应名字
 if self._drawing:
  label_color = self.label_colors[self._cur_label] if self._cur_label in self.label_colors else COLOR_GRAY
  if (self._pt1[0] >= self._pt0[0]) and (self._pt1[1] >= self._pt1[0]):
   cv2.rectangle(canvas, self._pt0, self._pt1, label_color, thickness=2)
  cv2.putText(canvas, self._cur_label, (self._pt0[0] + 3, self._pt0[1] + 15),
     cv2.FONT_HERSHEY_SIMPLEX, 0.5, label_color, 2)
 return canvas

# 利用repr()函数导出标注框数据到文件
@staticmethod
def export_bbox(filepath, bboxes):
 if bboxes:
  with open(filepath, 'w') as f:
   for bbox in bboxes:
    line = repr(bbox) + '\n'
    f.write(line)
 elif os.path.exists(filepath):
  os.remove(filepath)

# 利用eval()函数读取标注框字符串到数据
@staticmethod
def load_bbox(filepath):
 bboxes = []
 with open(filepath, 'r') as f:
  line = f.readline().rstrip()
  while line:
   bboxes.append(eval(line))
   line = f.readline().rstrip()
 return bboxes

# 利用eval()函数读取物体及对应颜色信息到数据
@staticmethod
def load_labels(filepath):
 label_colors = {}
 with open(filepath, 'r') as f:
  line = f.readline().rstrip()
  while line:
   label, color = eval(line)
   label_colors[label] = color
   line = f.readline().rstrip()
 print label_colors
 return label_colors

# 读取图像文件和对应标注框信息（如果有的话）
@staticmethod
def load_sample(filepath):
 img = cv2.imread(filepath)
 bbox_filepath = get_bbox_name(filepath)
 bboxes = []
 if os.path.exists(bbox_filepath):
  bboxes = SimpleBBoxLabeling.load_bbox(bbox_filepath)
 return img, bboxes

# 导出当前标注框信息并清空
def _export_n_clean_bbox(self):
 bbox_filepath = os.sep.join([self._data_dir, get_bbox_name(self._filelist[self._index])])
 self.export_bbox(bbox_filepath, self._bboxes)
 self._clean_bbox()

# 删除当前样本和对应的标注框信息
def _delete_current_sample(self):
 filename = self._filelist[self._index]
 filepath = os.sep.join([self._data_dir, filename])
 if os.path.exists(filepath):
   os.remove(filepath)
 filepath = get_bbox_name(filepath)
 if os.path.exists(filepath):
   os.remove(filepath)
 self._filelist.pop(self._index)
 print('{} is deleted!'.format(filename))

# 开始OpenCV窗口循环的方法，程序的主逻辑
def start(self):
 # 之前标注的文件名，用于程序判断是否需要执行一次图像读取
 last_filename = ''

# 标注物体在列表中的下标
 label_index = 0

# 所有标注物体名称的列表
 labels = self.label_colors.keys()

# 带标注物体的种类数
 n_labels = len(labels)

# 定义窗口和鼠标回调
 cv2.namedWindow(self.window_name)
 cv2.setMouseCallback(self.window_name, self._mouse_ops)
 key = KEY_EMPTY

# 定义每次循环的持续时间
 delay = int(1000 / FPS)

# 只要没有按下Delete键，就持续循环
 while key != KEY_DELETE:
  # 上下方向键选择当前标注物体
  if key == KEY_UP:
   if label_index == 0:
    pass
   else:
    label_index -= 1
  elif key == KEY_DOWN:
   if label_index == n_labels - 1:
    pass
   else:
    label_index += 1
  # 左右方向键选择标注图片
  elif key == KEY_LEFT:
   # 已经到了第一张图片的话就不需要清空上一张
   if self._index > 0:
    self._export_n_clean_bbox()
   self._index -= 1
   if self._index < 0:
    self._index = 0
  elif key == KEY_RIGHT:
   # 已经到了最后一张图片的就不需要清空上一张
   if self._index < len(self._filelist) - 1:
    self._export_n_clean_bbox()
   self._index += 1
   if self._index > len(self._filelist) - 1:
    self._index = len(self._filelist) - 1
  # 删除当前图片和对应标注的信息
  elif key == KEY_DELETE:
   if askyesno('Delete Sample', 'Are you sure?'):
    self._delete_current_sample()
    key = KEY_EMPTY
    continue
  # 如果键盘操作执行了换图片， 则重新读取， 更新图片
  filename = self._filelist[self._index]
  if filename != last_filename:
   filepath = os.sep.join([self._data_dir, filename])
   img, self._bboxes = self.load_sample(filepath)
  # 更新当前标注物体名称
  self._cur_label = labels[label_index]
  # 把标注和相关信息画在图片上并显示指定的时间
  canvas = self._draw_bbox(img)
  cv2.imshow(self.window_name, canvas)
  key = cv2.waitKey(delay)
  # 当前文件名就是下次循环的老文件名
  last_filename = filename
 print 'Finished!'
 cv2.destroyAllWindows()
 #如果退出程序，需要对当前文件进行保存
 self.export_bbox(os.sep.join([self._data_dir, get_bbox_name(filename)]), self._bboxes)
 print 'Labels updated!'

以上实现了工具类，当然需要一个入口函数，将工具类保存为SimpleBBoxLabeling.py，新建Run_Detect.py，写以下内容：

# coding:utf-8

# tkinter是Python内置的简单GUI库，实现打开文件夹、确认删除等操作十分方便
from tkFileDialog import askdirectory
# 导入创建的工具类
from SimpleBBoxLabeling import SimpleBBoxLabeling

if __name__ == '__main__':
dir_with_images = askdirectory(title='Where is the images?')
labeling_task = SimpleBBoxLabeling(dir_with_images)
labeling_task.start()

以下是实现后的效果：

需要的文件

.labels文件内容格式

选择文件夹

进行标注

生成相应标签内容

标注结果

标注后的文件格式为：物体，左上角（起点）和右下角（终点）的坐标。

参考资料： 《深度学习与计算机视觉——算法原理、框架应用与代码实现》 叶韵（编著）

来源：https://blog.csdn.net/weixin_34290352/article/details/87630162