针对目标检测标注数据做的数据预处理操作。

    1. # -*- coding: utf-8 -*-
    2. """
    3.  @File    : rotate_images.py
    4.  @Time    : 2020-8-19 21:21
    5.  @Author  : qian733
    6.  @Description    :
    7. """
    8. import os
    9. import cv2
    10. import shutil
    11. import argparse
    12. import voc_utils
    13. import numpy as np
    16. class RotateProcessing(object):
    17.     def __init__(self, angle, center=0, scale=1.0):
    18.         """
    20.         :param angle: 旋转角度
    21.         :param center: 旋转的中心点，默认为图像中心，以0代替
    22.         :param scale: 图像缩放因子
    23.         """
    24.         self.angle = angle
    25.         self.center = center
    26.         self.scale = scale
    28.     def rotate(self, image):
    29.         (h, w) = image.shape[:2]
    30.         if not self.center:
    31.             (cx, cy) = (w / 2, h / 2)
    32.         else:
    33.             (cx, cy) = self.center  # 二元元组
    35.         # 设置旋转矩阵
    36.         matrix = cv2.getRotationMatrix2D((cx, cy), -self.angle, self.scale)
    37.         cos = np.abs(matrix[0, 0])
    38.         sin = np.abs(matrix[0, 1])
    40.         # 计算图像旋转后的新边界
    41.         new_weight = int((h * sin) + (w * cos))
    42.         new_height = int((h * cos) + (w * sin))
    44.         # 调整旋转矩阵的移动距离（t_{x}, t_{y}）
    45.         matrix[0, 2] += (new_weight / 2) - cx
    46.         matrix[1, 2] += (new_height / 2) - cy
    48.         return cv2.warpAffine(image, matrix, (new_weight, new_height))
    50.     def get_rotate_boxes(self, img_shape, gt_bboxes):
    51.         """
    53.         :param img_shape:
    54.         :param gt_bboxes: 由TensorFlow返回的预测结果
    55.         :return:
    56.         """
    57.         num_boxes = len(gt_bboxes)
    58.         (h, w) = img_shape[0], img_shape[1]
    59.         mask = np.zeros(shape=(h, w, num_boxes), dtype=np.uint8)
    61.         for i, box in enumerate(gt_bboxes):
    62.             y1, x1, y2, x2 = box
    63.             # 这里的gt是大于0的float，不需要分别乘以长宽，但是需要取整
    64.             y1, y2 = int(np.floor(y1)), int(np.ceil(y2))
    65.             x1, x2 = int(np.floor(x1)), int(np.ceil(x2))
    66.             mask[y1:y2, x1:x2, i] = 1
    68.         if not self.center:
    69.            (cx, cy) = (w / 2, h / 2)
    70.         else:
    71.             (cx, cy) = self.center  # 二元元组
    73.         matrix = cv2.getRotationMatrix2D((cx, cy), -self.angle, self.scale)
    74.         cos = np.abs(matrix[0, 0])
    75.         sin = np.abs(matrix[0, 1])
    76.         new_weight = int((h * sin) + (w * cos))
    77.         new_height = int((h * cos) + (w * sin))
    78.         matrix[0, 2] += (new_weight / 2) - cx
    79.         matrix[1, 2] += (new_height / 2) - cy
    80.         new_mask = cv2.warpAffine(mask, matrix, (new_weight, new_height))
    82.         if num_boxes == 1:
    83.             new_mask = new_mask[:, :, np.newaxis]
    84.         new_gt_boxes = np.zeros_like(gt_bboxes)
    85.         for i in range(num_boxes):
    86.             hs, ws = np.where(new_mask[:, :, i] == 1)
    87.             y1, x1, y2, x2 = 0, 0, 0, 0
    88.             if len(hs) > 0:
    89.                 y1, y2 = np.min(hs), np.max(hs)
    90.             if len(ws) > 0:
    91.                 x1, x2 = np.min(ws), np.max(ws)
    92.             new_gt_boxes[i] = np.array([y1, x1, y2, x2])
    93.         return new_gt_boxes, new_height, new_weight
    95. def rotate_images_special(args, out_anno, out_jpg, out_bg):
    96.     """
    98.     :param args:
    99.     :param out_anno:
    100.     :param out_jpg:
    101.     :param out_bg:
    102.     :return:
    103.     """
    104.     if args.subfile:
    105.         subfile_list = args.subfile.split("+")
    106.     else:
    107.         subfile_list = os.listdir(args.input_path)
    108.     angle_list = args.angle.split("+")
    110.     for sub in os.listdir(args.input_path):
    111.         if sub in subfile_list:
    112.             print(sub)
    113.             sub_dir = os.path.join(args.input_path, sub)  # 如 root_dir/small_0511 下面放着图像和xml文件
    114.             anno_dir = sub_dir + "/Annotations/"
    115.             img_dir = sub_dir + "/JPEGImages/"
    116.             bg_list = []
    117.             for anno in os.listdir(anno_dir):
    118.                 xml_path = os.path.join(anno_dir, anno)  # xml
    119.                 _, height, width, class_names, gt_boxes = voc_utils.parse_voc(xml_path)
    120.                 # 如果class为空，则进行去重；如果class只有第一类，则进行旋转
    121.                 if args.exclude_class not in class_names:
    122.                     for angle in angle_list:
    123.                         RP = RotateProcessing(int(angle))
    124.                         image_ori = cv2.imread(img_dir + anno.split(".")[0] + ".jpg")
    125.                         image_rotate = RP.rotate(image_ori)
    126.                         cv2.imwrite(out_jpg + "/r" + angle + anno.split(".")[0] + ".jpg", image_rotate)
    127.                         image_rotate_boxes, nH, nW = RP.get_rotate_boxes(image_ori.shape, gt_boxes)
    128.                         voc_utils.save_to_xml_file("r" + angle + anno.split(".")[0] + ".jpg", nH, nW, class_names, image_rotate_boxes, out_anno)
    130. def rotate_images(args, out_anno, out_jpg, out_bg):
    131.     """
    133.     :param args:
    134.     :param out_anno:
    135.     :param out_jpg:
    136.     :param out_bg:
    137.     :return:
    138.     """
    139.     if args.subfile:
    140.         subfile_list = args.subfile.split("+")
    141.     else:
    142.         subfile_list = os.listdir(args.input_path)
    143.     angle_list = args.angle.split("+")
    145.     for sub in os.listdir(args.input_path):
    146.         if sub in subfile_list:
    147.             print(sub)
    148.             sub_dir = os.path.join(args.input_path, sub)  # 如 root_dir/small_0511 下面放着图像和xml文件
    149.             anno_dir = sub_dir + "/Annotations/"
    150.             img_dir = sub_dir + "/JPEGImages/"
    151.             bg_list = []
    152.             for anno in os.listdir(anno_dir):
    153.                 xml_path = os.path.join(anno_dir, anno)  # xml
    154.                 _, height, width, class_names, gt_boxes = voc_utils.parse_voc(xml_path)
    155.                 # 如果class为空，则旋转图像并修改xml中对应的图像宽高
    156.                 if not class_names:
    157.                                         for angle in angle_list:
    158.                         RP = RotateProcessing(int(angle))
    159.                         image_ori = cv2.imread(img_dir + anno.split(".")[0] + ".jpg")
    160.                         image_rotate = RP.rotate(image_ori)
    161.                         cv2.imwrite(out_jpg + "/r" + anno.split(".")[0]+ "_" + angle + ".jpg", image_rotate)
    162.                         print("touch " + anno.split(".")[0]+ "_" + angle + ".xml")
    164.                 elif args.exclude_class not in class_names:
    165.                     for angle in angle_list:
    166.                         RP = RotateProcessing(int(angle))
    167.                         image_ori = cv2.imread(img_dir + anno.split(".")[0] + ".jpg")
    168.                         image_rotate = RP.rotate(image_ori)
    169.                         cv2.imwrite(out_jpg + "/r" + anno.split(".")[0]+ "_" + angle + ".jpg", image_rotate)
    170.                         image_rotate_boxes, nH, nW = RP.get_rotate_boxes(image_ori.shape, gt_boxes)
    171.                         voc_utils.save_to_xml_file("r" + anno.split(".")[0]+ "_" + angle + ".jpg", nH, nW, class_names,
    172.                                                    image_rotate_boxes, out_anno)
    175. def main():
    176.     parse = argparse.ArgumentParser()
    177.     parse.add_argument("--input_path", type=str,
    178.                        default="/sdb/tmp/Behavior_Normalization_Monitoring/data/Tagged_data_crop/small")
    179.     parse.add_argument("--out_path", type=str,
    180.                        default="/sdb/tmp/Behavior_Normalization_Monitoring/data/Tagged_data_rotate/small")
    181.     parse.add_argument("--angle", type=str, default="90+180+270")
    182.     parse.add_argument("--subfile", type=str, default="")
    183.     parse.add_argument("--exclude_class", type=str, default="phone")
    184.     arg = parse.parse_args()
    186.     out_anno = arg.out_path + "/Annotations/"
    187.     out_jpg = arg.out_path + "/JPEGImages/"
    188.     out_bg = arg.out_path + "/bg/"
    189.     if not os.path.exists(arg.out_path):
    190.         os.makedirs(out_anno)
    191.         os.makedirs(out_jpg)
    192.         os.makedirs(out_bg)
    194.     rotate_images(arg, out_anno, out_jpg, out_bg)
    197. if __name__ == '__main__':
    198.     main()