devbox
繁依Fanyi0
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

YOLO交通标志识别数据集准备_交通识别数据

作者：繁依Fanyi0 | 2024-07-16 03:20:49

赞

踩

交通识别数据

index (tsinghua.edu.cn)

  

 

模板转换jinjia2包链接：https://pan.baidu.com/s/1ycb_zf8oS88HF0FvpXrYFg?pwd=ym9n 
提取码：ym9n 

  1. import os
  2. from jinja2 import Environment, PackageLoader
  3.  
  4.  
  5. class xml_fill:
  6.     def __init__(self, path, width, height, depth=3, database='Unknown', segmented=0):
  7.         environment = Environment(loader=PackageLoader('source', 'XML_template'), keep_trailing_newline=True)
  8.         self.annotation_template = environment.get_template('voc_template.xml')
  9.  
  10.         abspath = os.path.abspath(path)
  11.  
  12.         self.template_parameters = {
  13.             'path': abspath,
  14.             'filename': os.path.basename(abspath),
  15.             'folder': os.path.basename(os.path.dirname(abspath)),
  16.             'width': width,
  17.             'height': height,
  18.             'depth': depth,
  19.             'database': database,
  20.             'segmented': segmented,
  21.             'objects': []
  22.         }
  23.  
  24.     def add_obj_box(self, name, xmin, ymin, xmax, ymax, pose='Unspecified', truncated=0, difficult=0):
  25.         self.template_parameters['objects'].append({
  26.             'name': name,
  27.             'xmin': xmin,
  28.             'ymin': ymin,
  29.             'xmax': xmax,
  30.             'ymax': ymax,
  31.             'pose': pose,
  32.             'truncated': truncated,
  33.             'difficult': difficult,
  34.         })
  35.  
  36.     def save_xml(self, annotation_path):
  37.         with open(annotation_path, 'w') as file:
  38.             content = self.annotation_template.render(**self.template_parameters)
  39.             file.write(content)

 

  1. import json
  2. import os
  3. from PIL import Image
  4. from voc_xml_generator import xml_fill
  5. tt100k_parent_dir = "G:\\"
  6.  
  7. def find_image_size(filename):
  8.     with Image.open(filename) as img:
  9.         img_weight = img.size[0]
  10.         img_height = img.size[1]
  11.         img_depth = 3
  12.         return img_weight, img_height, img_depth
  13.  
  14. def load_mask(annos, datadir, imgid, filler):
  15.     img = annos["imgs"][imgid]
  16.     path = img['path']
  17.     for obj in img['objects']:
  18.         name = obj['category']
  19.         box = obj['bbox']
  20.         xmin = int(box['xmin'])
  21.         ymin = int(box['ymin'])
  22.         xmax = int(box['xmax'])
  23.         ymax = int(box['ymax'])
  24.         filler.add_obj_box(name, xmin, ymin, xmax, ymax)
  25.  
  26. work_sapce_dir = os.path.join(tt100k_parent_dir, "TT100K\\VOCdevkit\\")
  27. if not os.path.isdir(work_sapce_dir):
  28.     os.mkdir(work_sapce_dir)
  29. work_sapce_dir = os.path.join(work_sapce_dir, "VOC20230102\\")
  30. if not os.path.isdir(work_sapce_dir):
  31.     os.mkdir(work_sapce_dir)
  32. jpeg_images_path = os.path.join(work_sapce_dir, 'JPEGImages')
  33. annotations_path = os.path.join(work_sapce_dir, 'Annotations')
  34. if not os.path.isdir(jpeg_images_path):
  35.     os.mkdir(jpeg_images_path)
  36. if not os.path.isdir(annotations_path):
  37.         os.mkdir(annotations_path)
  38.  
  39. datadir = tt100k_parent_dir + "TT100K\\data"
  40.  
  41. filedir = datadir + "\\annotations.json"
  42. ids = open(datadir + "\\train\\ids.txt").read().splitlines()
  43. annos = json.loads(open(filedir).read())
  44.  
  45.  
  46. for i, value in enumerate(ids):
  47.     imgid = value
  48.     filename = datadir + "\\train\\" + imgid + ".jpg"
  49.     width,height,depth = find_image_size(filename)
  50.     filler = xml_fill(filename, width, height, depth)
  51.     load_mask(annos, datadir, imgid, filler)
  52.     filler.save_xml(annotations_path + '\\' + imgid + '.xml')
  53.     print("%s.xml saved\n"%imgid)
  54.

 

  1. import xml.etree.ElementTree as ET
  2. import os
  3. import random
  4. from shutil import move
  5.  
  6. type45="i2,i4,i5,il100,il60,il80,io,ip,p10,p11,p12,p19,p23,p26,p27,p3,p5,p6,pg,ph4,ph4.5,ph5,pl100,pl120,pl20,pl30,pl40,pl5,pl50,pl60,pl70,pl80,pm20,pm30,pm55,pn,pne,po,pr40,w13,w32,w55,w57,w59,wo"
  7. type45 = type45.split(',')
  8. classes = type45
  9.  
  10. TRAIN_RATIO = 80
  11.  
  12. def clear_hidden_files(path):
  13.     dir_list = os.listdir(path)
  14.     for i in dir_list:
  15.         abspath = os.path.join(os.path.abspath(path), i)
  16.         if os.path.isfile(abspath):
  17.             if i.startswith("._"):
  18.                 os.remove(abspath)
  19.         else:
  20.             clear_hidden_files(abspath)
  21.  
  22. def convert(size, box):
  23.     dw = 1./size[0]
  24.     dh = 1./size[1]
  25.     x = (box[0] + box[1])/2.0
  26.     y = (box[2] + box[3])/2.0
  27.     w = box[1] - box[0]
  28.     h = box[3] - box[2]
  29.     x = x*dw
  30.     w = w*dw
  31.     y = y*dh
  32.     h = h*dh
  33.     return (x,y,w,h)
  34.  
  35. def convert_annotation(image_id):
  36.     in_file = open('VOC/2022/ANNOTATIONS/%s.xml' %image_id)
  37.     out_file = open('VOC/2022/YOLOLabels/%s.txt' %image_id, 'w')
  38.     tree=ET.parse(in_file)
  39.     root = tree.getroot()
  40.     size = root.find('size')
  41.     w = int(size.find('width').text)
  42.     h = int(size.find('height').text)
  43.  
  44.     for obj in root.iter('object'):
  45.         difficult = obj.find('difficult').text
  46.         cls = obj.find('name').text
  47.         if cls not in classes or int(difficult) == 1:
  48.             continue
  49.         cls_id = classes.index(cls)
  50.         xmlbox = obj.find('bndbox')
  51.         b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
  52.         bb = convert((w,h), b)
  53.         out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
  54.     in_file.close()
  55.     out_file.close()
  56.  
  57. wd = os.getcwd()
  58. wd = os.getcwd()
  59. data_base_dir = os.path.join(wd, "VOC/")
  60. if not os.path.isdir(data_base_dir):
  61.     os.mkdir(data_base_dir)
  62. work_sapce_dir = os.path.join(data_base_dir, "2022/")
  63. if not os.path.isdir(work_sapce_dir):
  64.     os.mkdir(work_sapce_dir)
  65. annotation_dir = os.path.join(work_sapce_dir, "ANNOTATIONS/")
  66. if not os.path.isdir(annotation_dir):
  67.         os.mkdir(annotation_dir)
  68. clear_hidden_files(annotation_dir)
  69. image_dir = os.path.join(work_sapce_dir, "IMAGE/")
  70. if not os.path.isdir(image_dir):
  71.         os.mkdir(image_dir)
  72. clear_hidden_files(image_dir)
  73. yolo_labels_dir = os.path.join(work_sapce_dir, "YOLOLabels/")
  74. if not os.path.isdir(yolo_labels_dir):
  75.         os.mkdir(yolo_labels_dir)
  76. clear_hidden_files(yolo_labels_dir)
  77. yolov5_images_dir = os.path.join(data_base_dir, "images/")
  78. if not os.path.isdir(yolov5_images_dir):
  79.         os.mkdir(yolov5_images_dir)
  80. clear_hidden_files(yolov5_images_dir)
  81. yolov5_labels_dir = os.path.join(data_base_dir, "labels/")
  82. if not os.path.isdir(yolov5_labels_dir):
  83.         os.mkdir(yolov5_labels_dir)
  84. clear_hidden_files(yolov5_labels_dir)
  85. yolov5_images_train_dir = os.path.join(yolov5_images_dir, "train/")
  86. if not os.path.isdir(yolov5_images_train_dir):
  87.         os.mkdir(yolov5_images_train_dir)
  88. clear_hidden_files(yolov5_images_train_dir)
  89. yolov5_images_test_dir = os.path.join(yolov5_images_dir, "val/")
  90. if not os.path.isdir(yolov5_images_test_dir):
  91.         os.mkdir(yolov5_images_test_dir)
  92. clear_hidden_files(yolov5_images_test_dir)
  93. yolov5_labels_train_dir = os.path.join(yolov5_labels_dir, "train/")
  94. if not os.path.isdir(yolov5_labels_train_dir):
  95.         os.mkdir(yolov5_labels_train_dir)
  96. clear_hidden_files(yolov5_labels_train_dir)
  97. yolov5_labels_test_dir = os.path.join(yolov5_labels_dir, "val/")
  98. if not os.path.isdir(yolov5_labels_test_dir):
  99.         os.mkdir(yolov5_labels_test_dir)
  100. clear_hidden_files(yolov5_labels_test_dir)
  101.  
  102. train_file = open(os.path.join(wd, "yolov5_train.txt"), 'w')
  103. test_file = open(os.path.join(wd, "yolov5_val.txt"), 'w')
  104. train_file.close()
  105. test_file.close()
  106. train_file = open(os.path.join(wd, "yolov5_train.txt"), 'a')
  107. test_file = open(os.path.join(wd, "yolov5_val.txt"), 'a')
  108. list_imgs = os.listdir(image_dir) # list image files
  109. prob = random.randint(1, 100)
  110. for i in range(0,len(list_imgs)):
  111.     path = os.path.join(image_dir,list_imgs[i])
  112.     if os.path.isfile(path):
  113.         image_path = image_dir + list_imgs[i]
  114.         voc_path = list_imgs[i]
  115.         (nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
  116.         (voc_nameWithoutExtention, voc_extention) = os.path.splitext(os.path.basename(voc_path))
  117.         annotation_name = nameWithoutExtention + '.xml'
  118.         annotation_path = os.path.join(annotation_dir, annotation_name)
  119.         label_name = nameWithoutExtention + '.txt'
  120.         label_path = os.path.join(yolo_labels_dir, label_name)
  121.     prob = random.randint(1, 100)
  122.     if(prob < TRAIN_RATIO): # train dataset
  123.         if os.path.exists(annotation_path):
  124.             train_file.write(image_path + '\n')
  125.             convert_annotation(nameWithoutExtention) # convert label
  126.             move(image_path, yolov5_images_train_dir + voc_path)
  127.             move(label_path, yolov5_labels_train_dir + label_name)
  128.     else: # test dataset
  129.         if os.path.exists(annotation_path):
  130.             test_file.write(image_path + '\n')
  131.             convert_annotation(nameWithoutExtention) # convert label
  132.             move(image_path, yolov5_images_test_dir + voc_path)
  133.             move(label_path, yolov5_labels_test_dir + label_name)
  134. train_file.close()
  135. test_file.close()

 

 

  1. class LoadImgLabels(Dataset):
  2.     # root = "YOLO/VOC"
  3.     def __init__(self,root,mode):
  4.         super(LoadImgLabels, self).__init__()
  5.         self.root = root
  6.         self.mode = mode
  7.  
  8.         img_path = get_path(os.path.join(root,'images',self.mode))
  9.         lab_path = get_path(os.path.join(root,'labels',self.mode))
  10.         self.img_files = get_file(img_path)
  11.         self.label_files = img2label_paths(self.img_files)
  12.  
  13.  
  14.     def __len__(self):
  15.         return len()
  16.  
  17.     def __getitem__(self, item):
  18.         return

# 获得（不同操作系统）标准路径
  1. def get_path(path):
  2.     p = str(Path(path))
  3.     return p
# 得到路径下的每个文件
  1. def get_file(path):
  2.     file = []
  3.     if os.path.isdir(path):
  4.         file += glob.iglob(path + os.sep + '*.*')
  5.     return file
# 由图片的文件得到对应标签的文件
  1. def img2label_paths(img_paths):
  2.     sa, sb = os.sep + 'images' + os.sep, os.sep + 'labels' + os.sep
  3.     return [x.replace(sa, sb, 1).replace(os.path.splitext(x)[-1], '.txt') for x in img_paths]
# 缓存标签
  1. def cache_labels(img_files, label_files, path='labels.cache'):
  2.     x = {}
  3.     pbar = tqdm(zip(img_files, label_files), desc='Scanning images', total=len(img_files))
  4.     for (img, label) in pbar:
  5.         print(img,label)
  6.         try:
  7.             l=[]
  8.             im = Image.open(img)
  9.             im.verify()
  10.             shape = im.size
  11.             if os.path.isfile(label):
  12.                 with open(label,'r') as f:
  13.                     l = np.array([x.split() for x in f.read().splitlines()], dtype=np.float32)
  14.             if len(l) == 0:
  15.                 l = np.zeros((0, 5), dtype=np.float32)
  16.             x[img] = [l,shape]
  17.         except:
  18.             pass
  19.     torch.save(x, path)
  20.     return x

  1. class LoadImgLabels(Dataset):
  2.     # root = "../VOC"
  3.     def __init__(self,root,mode,img_size):
  4.         super(LoadImgLabels, self).__init__()
  5.         self.root = root
  6.         self.mode = mode
  7.         self.img_size = img_size  # 输入图片分辨率大小
  8.  
  9.         img_path = get_path(os.path.join(root,'images',self.mode))
  10.  
  11.         if os.path.isfile('labels.cache'):
  12.             print("读取缓存标签文件'labels.cache'")
  13.             cache = torch.load('labels.cache')
  14.         else:
  15.             print("生成缓存标签文件'labels.cache'")
  16.             self.img_files = get_file(img_path)
  17.             self.label_files = img2label_paths(self.img_files)
  18.             cache = cache_labels(self.img_files, self.label_files)
  19.  
  20.         labels, shapes = zip(*cache.values())
  21.         self.labels = list(labels)
  22.         self.shapes = np.array(shapes, dtype=np.float64)
  23.         self.img_files = list(cache.keys())
  24.         self.label_files = img2label_paths(cache.keys())
  25.  
  26.  
  27.     def __len__(self):
  28.         return len(self.img_files)
  29.  
  30.     def __getitem__(self, index):
  31.  
  32.         return 0


# 加载图片

并根据设定的输入大小与图片原大小的比例ratio进行resize;

if img_size = 640:(1080, 1920)———>(360, 640)

  1. def load_image(img_files, img_size , index):    #   img_size = 640
  2.     path = img_files[index]
  3.     img = cv2.imread(path)
  4.     h0 ,w0 = img.shape[:2]
  5.     r = img_size / max(h0,w0)
  6.     if r != 1:
  7.         interp = cv2.INTER_AREA if r < 1 else cv2.INTER_LINEAR
  8.         img = cv2.resize(img, (int(w0 * r), int(h0 * r)), interpolation=interp)
  9.     return img, (h0, w0), img.shape[:2] #   (1080, 1920)———>(360, 640)
# 图像缩放: 保持图片的宽高比例，剩下的部分采用灰色填充。
  1. def Make_squqre(img, new_shape=(640, 640), color=(114, 114, 114)):
  2.     # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232
  3.     shape = img.shape[:2]  # 当前图片大小
  4.     if isinstance(new_shape, int):
  5.         new_shape = (new_shape, new_shape)
  6.  
  7.     # ----------------计算填充大小-----------------------------------------
  8.     r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])# r = 1.0
  9.     ratio = r, r  # ratio = (1.0,1.0)
  10.     new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
  11.     dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # 填充宽度,高度
  12.     # 计算上下左右填充大小
  13.     dw /= 2
  14.     dh /= 2
  15.     top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
  16.     left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
  17.     # ------------------进行填充-------------------------------------------
  18.     img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
  19.     return img, ratio, (dw, dh)

对标签的处理

# 1根据pad调整框的标签坐标

# 2调整框的标签，xyxy->xywh

# 3归一化标签0 - 1

  1.         labels = []
  2.         x = self.labels[index]
  3.         if x.size > 0:
  4.             # 根据pad调整框的标签坐标:注意label是真实位置,没有归一化的
  5.             labels = x.copy()
  6.             labels[:, 1] = ratio[0] * w * (x[:, 1] - x[:, 3] / 2) + pad[0]
  7.             labels[:, 2] = ratio[1] * h * (x[:, 2] - x[:, 4] / 2) + pad[1]
  8.             labels[:, 3] = ratio[0] * w * (x[:, 1] + x[:, 3] / 2) + pad[0]
  9.             labels[:, 4] = ratio[1] * h * (x[:, 2] + x[:, 4] / 2) + pad[1]
  10.         nL = len(labels)
  11.         if nL:  
  12.             labels[:, 1:5] = xyxy2xywh(labels[:, 1:5])
  13.             # 重新归一化标签0 - 1
  14.             labels[:, [2, 4]] /= img.shape[0]  # normalized height 0~1
  15.             labels[:, [1, 3]] /= img.shape[1]  # normalized width 0~1
  16.         labels_out = torch.zeros((nL, 6))
  17.         if nL:
  18.             labels_out[:, 1:] = torch.from_numpy(labels)
# 左上角右下角坐标格式转换成中心点+宽高坐标格式
  1. def xyxy2xywh(x):
  2.     # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
  3.     y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x)
  4.     y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center
  5.     y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
  6.     y[:, 2] = x[:, 2] - x[:, 0]  # width
  7.     y[:, 3] = x[:, 3] - x[:, 1]  # height
  8.     return y

dataloader

  1. import torch
  2. from contextlib import contextmanager
  3. from tqdm import tqdm
  4.  
  5. from YOLO.dataset.dataset import LoadImgLabels
  6.  
  7.  
  8.  
  9.  
  10.  
  11. # 定义生成器 _RepeatSampler
  12. class _RepeatSampler(object):
  13.     def __init__(self, sampler):
  14.         self.sampler = sampler
  15.  
  16.     def __iter__(self):
  17.         while True:
  18.             yield from iter(self.sampler)
  19. # 定义DataLoader（一个python生成器）
  20. class InfiniteDataLoader(torch.utils.data.dataloader.DataLoader):
  21.     def __init__(self, *args, **kwargs):
  22.         super().__init__(*args, **kwargs)
  23.         object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))
  24.         self.iterator = super().__iter__()
  25.  
  26.     def __len__(self):
  27.         return len(self.batch_sampler.sampler)
  28.  
  29.     def __iter__(self): # 实现了__iter__方法的对象是可迭代的
  30.         for i in range(len(self)):
  31.             yield next(self.iterator)
  32. @contextmanager
  33. def torch_distributed_zero_first(local_rank: int):""
  34.     if local_rank not in [-1, 0]:
  35.         torch.distributed.barrier() # Synchronizes all processes
  36.     yield
  37.     if local_rank == 0:
  38.         torch.distributed.barrier()
  39.  
  40.  
  41.  
  42. # 利用自定义的数据集(LoadImagesAndLabels)创建dataloader
  43. def create_dataloader(path, mode , imgsz, batch_size,rank=-1):
  44.     with torch_distributed_zero_first(rank):
  45.         dataset = LoadImgLabels(path, mode, imgsz)
  46.  
  47.     batch_size = min(batch_size, len(dataset))
  48.     dataloader = InfiniteDataLoader(dataset,#  torch.utils.data.DataLoader
  49.                                     batch_size=batch_size,
  50.                                     shuffle=True,
  51.                                     collate_fn=LoadImgLabels.collate_fn,
  52.                                     pin_memory=True)
  53.     return dataloader, dataset
  54.  
  55. dataloader, dataset = create_dataloader("G:\VOC", 'train',640, 2)
  56.  
  57. pbar = enumerate(dataloader)
  58. nb = len(dataloader)
  59. pbar = tqdm(pbar, total=nb)
  60.  
  61. for i, (imgs, targets, path) in pbar:
  62.     ni = i + nb * 1
  63.     imgs = imgs / 255.0
  64.     print(imgs.size(),targets.size())

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/繁依Fanyi0/article/detail/832116
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号