CCPD数据集处理

admin

[mw_shl_code=applescript,true]"""
@Author: HuKai
@Date: 2022/5/29  10:44
@github: https://github.com/HuKai97
"""
import os
import cv2
from shutil import copy2

def txt_translate(images_dir,filename, labels_dir):
    list1 = filename.split("-", 3)  # 第一次分割，以减号'-'做分割
    subname = list1[2]
    list2 = filename.split(".", 1)
    subname1 = list2[1]
    if subname1 == 'txt':
        return
    lt, rb = subname.split("_", 1)  # 第二次分割，以下划线'_'做分割
    lx, ly = lt.split("&", 1)
    rx, ry = rb.split("&", 1)
    width = int(rx) - int(lx)
    height = int(ry) - int(ly)  # bounding box的宽和高
    cx = float(lx) + width / 2
    cy = float(ly) + height / 2  # bounding box中心点

    img = cv2.imread(images_dir + filename)
    if img is None:  # 自动删除失效图片（下载过程有的图片会存在无法读取的情况）
        os.remove(images_dir + filename)
        return
    width = width / img.shape[1]
    height = height / img.shape[0]
    cx = cx / img.shape[1]
    cy = cy / img.shape[0]

    txtname = filename.split(".", 1)
    txtfile = labels_dir + txtname[0] + ".txt"
        # 蓝牌是第0类，绿牌是第1类
    with open(txtfile, "w") as f:
        f.write(str(1) + " " + str(cx) + " " + str(cy) + " " + str(width) + " " + str(height))

images_trainDir = "../../CCPD2019/images/train/"#（将图片文件夹中的6份放在这个文件夹下）
images_validDir = "../../CCPD2019/images/val/"     #（将图片文件夹中的2份放在这个文件夹下）
images_detectDir = "../../CCPD2019/images/test/"   #（将图片文件夹中的2份放在这个文件夹下）
labels_trainDir = "../../CCPD2019/labels/train/"
labels_validDir = "../../CCPD2019/labels/val/"
labels_detectDir = "../../CCPD2019/labels/test/"
src_dir = "../../CCPD2020/ccpd_green/"


num = 0
files_f = os.listdir(src_dir)
for files in files_f:
    for file in os.listdir(src_dir+files):
        temp = num%10
        temp_dir = src_dir+files+'/'+file
        if temp < 6 :
            copy2(temp_dir,images_trainDir)
            txt_translate(src_dir+files+'/',file,labels_trainDir)
        elif temp >= 6 and temp < 8 :
            copy2(temp_dir,images_validDir)
            txt_translate(src_dir+files+'/',file,labels_validDir)
        else :
            copy2(temp_dir,images_detectDir)
            txt_translate(src_dir+files+'/',file,labels_detectDir)
        num += 1
# trainfiles = os.listdir(r"K:\MyProject\datasets\ccpd\new\ccpd_2019\base")  #（图片文件夹）
# num_train = len(trainfiles)
# print("num_train: " + str(num_train) )
# index_list = list(range(num_train))
# print(index_list)
# random.shuffle(index_list)  # 打乱顺序
# num = 0
# trainDir = r"K:\MyProject\datasets\ccpd\new\ccpd_2019\train"   #（将图片文件夹中的6份放在这个文件夹下）
# validDir = r"K:\MyProject\datasets\ccpd\new\ccpd_2019\val"     #（将图片文件夹中的2份放在这个文件夹下）
# detectDir = r"K:\MyProject\datasets\ccpd\new\ccpd_2019\test"   #（将图片文件夹中的2份放在这个文件夹下）
# for i in index_list:
#     fileName = os.path.join(r"K:\MyProject\datasets\ccpd\new\ccpd_2019\base", trainfiles)  #（图片文件夹）+图片名=图片地址
#     if num < num_train*0.7:  # 7:1:2
#         print(str(fileName))
#         copy2(fileName, trainDir)
#     elif num < num_train*0.8:
#         print(str(fileName))
#         copy2(fileName, validDir)
#     else:
#         print(str(fileName))
#         copy2(fileName, detectDir)
#     num += 1
[/mw_shl_code]