import glob
import os
import numpy as np
import shutil
from PIL import Image

START_BOUNDING_BOX_ID = 0

def val_train(SHHB_img, SHHB_txt,pro_train,pro_valid):
train_img = ‘D:\\project\\PyTorch-YOLOv3-master\\data\\custom\\images\\SHHB_train\\’ #分离的训练图片路径
val_img = ‘D:\\project\\PyTorch-YOLOv3-master\\data\\custom\\images\\SHHB_valid\\’ #分离的验证集图片路径

# train_img = \'D:\\\\project\\\\zzzzzzz\\\\ag_image\\\\image\\\\SHHB_train\\\\\'  # 分离的训练图片路径# val_img = \'D:\\\\project\\\\zzzzzzz\\\\ag_image\\\\image\\\\SHHB_valid\\\\\'       #分离的验证集图片路径if os.path.exists(train_img):         #创建文件夹shutil.rmtree(train_img)os.makedirs(train_img)if os.path.exists(val_img):shutil.rmtree(val_img)os.makedirs(val_img)train_ratio = 0.9        #保留90%的图片为训练图片train_list = glob.glob(SHHB_img + \'/*.jpg\')train_list = np.sort(train_list)np.random.seed(100)np.random.shuffle(train_list)train_num = int(len(train_list) * train_ratio)img_list_train = train_list[:train_num]img_list_val = train_list[train_num:]f1 = open(\'SHHB_train.txt\',\'w\')             #把数据的路径写入到文件夹# f1 = open(\'D:\\\\project\\\\zzzzzzz\\\\ag_image\\\\SHHB_train.txt\',\'w\')             #把数据的路径写入到文件夹f = open(pro_train, \'w\')for img in img_list_train:f1.write(img + \'\\n\')shutil.copy(img, train_img)       #图片训练和验证数据分离n_img = img.split(\'G_\')n_img = n_img[1].split(\'.\')                   #n_img[0]为图片的序列号for line in open(SHHB_txt,\'r\'):n_txt = line.split(\'G_\')n_txt = n_txt[1].split(\'.\')                #n_txt[0]为图片的序列号if n_txt[0] == n_img[0]:f.write(line)f1.close()f.close()f2 = open(\'SHHB_valid.txt\',\'w\')f = open(pro_valid, \'w\')for img in img_list_val:f2.write(img + \'\\n\')shutil.copy(img,val_img)n_img = img.split(\'G_\')n_img = n_img[1].split(\'.\')  # n_img[0]为图片的序列号for line in open(SHHB_txt, \'r\'):n_txt = line.split(\'G_\')n_txt = n_txt[1].split(\'.\')  # n_txt[0]为图片的序列号if n_txt[0] == n_img[0]:f.write(line)f2.close()f.close()

def convert(pro_train,pro_valid):
train_txt = ‘D:\\project\\PyTorch-YOLOv3-master\\data\\custom\\labels\\SHHB_train\\’ #分离的训练txt文件
val_txt = ‘D:\\project\\PyTorch-YOLOv3-master\\data\\custom\\labels\\SHHB_valid\\’ #分离的验证txt文件
# train_txt = ‘D:\\project\\zzzzzzz\\ag_image\\lable\\SHHB_train\\’ #分离的训练txt文件
# val_txt = ‘D:\\project\\zzzzzzz\\ag_image\\lable\\SHHB_valid\\’ #分离的验证txt文件
if os.path.exists(train_txt):
shutil.rmtree(train_txt)
os.makedirs(train_txt)
if os.path.exists(val_txt):
shutil.rmtree(val_txt)
os.makedirs(val_txt)

# categories = pre_define_categories.copy()# bnd_id = START_BOUNDING_BOX_ID# all_categories = {}flag_continue = [1,2,3,4,5]            #设置跳出循环flag = 0#训练集for line in open(pro_train,\'r\'):linefeed = False                          #消除第一行换行line_feed = False              #如果上一个人头信息舍去，则不用换行line = line.strip(\'\\n\')n_txt = line.split(\'G_\')n_txt = n_txt[1].split(\'.jpg\')# n_txt = n_txt[1].split(\' \')txt_id = n_txt[0]               #txt_id获取txt文件序号# txt_id = n_txt[0].split(\'.\')               #txt_id获取txt文件序号# txt_id = txt_id[0]txt_content = n_txt[1].split(\' \')# n_txt.remove(n_txt[0])# txt_content = n_txttxt_content.remove(txt_content[0])# txt_content.strip(\'\\n\')print(txt_id)txt_content = list(map(int, txt_content))         #把字符转换成int类型txt_total = txt_content[0] * 5        #文本每一行数字的总数txt_content.remove(txt_content[0])  # 移除第一列的人头数f = open(train_txt + \'\\\\\' + \'IMG_\' + txt_id + \'.txt\', \'w\')        #写入与图片对应的txt文件for i in range(txt_total):if i % 5 == 0 and linefeed and line_feed:f.write(\'\\n\')if i % 5 == 0:if txt_content[i + 3] <= 5 or txt_content[i + 4] <= 5:flag = 5else:txt_content[i] == 0# elif (i % 5 == 1) and (txt_content[i + 2] <=5 or txt_content[i + 3] <= 5):#     flag = 4if flag in flag_continue:flag = flag - 1line_feed = Falsecontinueelse:                #把人头信息写入对应的txt文件中linefeed = Trueline_feed = Truefilename = \'D:\\\\project\\\\PyTorch-YOLOv3-master\\\\data\\\\custom\\\\images\\\\SHHB_train\\\\\' + \'IMG_\' + txt_id + \'.jpg\'# filename = \'D:\\\\project\\\\zzzzzzz\\\\ag_image\\\\image\\\\SHHB_train\\\\\' + \'IMG_\' + txt_id + \'.jpg\'img = Image.open(filename)           #获得图片的sizeif i % 5 == 1:f.write(str(((txt_content[i] * 2 + txt_content[i + 2]) / 2) / img.size[0]) + \' \')elif i % 5 == 2:f.write(str(((txt_content[i] * 2 + txt_content[i + 2]) / 2) / img.size[1]) + \' \')# elif i % 5 == 3:#     f.write(str(txt_content[i] / img.size[0]) + \' \')# elif i % 5 == 4:#     f.write(str(txt_content[i] / img.size[1]) + \' \')elif i % 5 == 3:f.write(str(txt_content[i] / ((img.size[0] + img.size[1]) / 2)) + \' \')elif i % 5 == 4:f.write(str(txt_content[i] / ((img.size[0] + img.size[1]) / 2)) + \' \')# if i % 5 == 1 or i % 5 == 3:#     f.write(str(txt_content[i] / img.size[0]) + \' \')# elif i % 5 == 2 or i % 5 == 4:#     f.write(str(txt_content[i] / img.size[1]) + \' \')else:f.write(str(txt_content[i]*0) + \' \')f.close()#验证集for line in open(pro_valid, \'r\'):linefeed = False  # 消除第一行换行line_feed = False  # 如果上一个人头信息舍去，则不用换行line = line.strip(\'\\n\')n_txt = line.split(\'G_\')n_txt = n_txt[1].split(\'.jpg\')txt_id = n_txt[0]  # txt_id获取txt文件序号txt_content = n_txt[1].split(\' \')txt_content.remove(txt_content[0])# txt_content.strip(\'\\n\')print(txt_id)txt_content = list(map(int, txt_content))txt_total = txt_content[0] * 5  # 文本每一行数字的总数txt_content.remove(txt_content[0])  # 移除第一列的人头数f = open(val_txt + \'\\\\\' + \'IMG_\' + txt_id + \'.txt\', \'w\')  # 写入与图片对应的txt文件for i in range(txt_total):if i % 5 == 0 and linefeed and line_feed:f.write(\'\\n\')if i % 5 == 0:if txt_content[i + 3] <= 5 or txt_content[i + 4] <= 5:flag = 5else:txt_content[i] == 0if flag in flag_continue:flag = flag - 1line_feed = Falsecontinueelse:  # 把人头信息写入对应的txt文件中linefeed = Trueline_feed = Truefilename = \'D:\\\\project\\\\PyTorch-YOLOv3-master\\\\data\\\\custom\\\\images\\\\SHHB_valid\\\\\' + \'IMG_\' + txt_id + \'.jpg\'img = Image.open(filename)  # 获得图片的sizeif i % 5 == 1:f.write(str(((txt_content[i] * 2 + txt_content[i + 2]) / 2) / img.size[0]) + \' \')elif i % 5 == 2:f.write(str(((txt_content[i] * 2 + txt_content[i + 2]) / 2) / img.size[1]) + \' \')# elif i % 5 == 3:#     f.write(str(txt_content[i] / img.size[0]) + \' \')# elif i % 5 == 4:#     f.write(str(txt_content[i] / img.size[1]) + \' \')elif i % 5 == 3:f.write(str(txt_content[i] / ((img.size[0] + img.size[1]) / 2)) + \' \')elif i % 5 == 4:f.write(str(txt_content[i] / ((img.size[0] + img.size[1]) / 2)) + \' \')# if i % 5 == 1 or i % 5 == 3:#     f.write(str(txt_content[i] / img.size[0]) + \' \')# elif i % 5 == 2 or i % 5 == 4:#     f.write(str(txt_content[i] / img.size[1]) + \' \')else:f.write(str(txt_content[i]*0) + \' \')f.close()

if name == ‘main’:
SHHB_img = ‘D:\\project\\crowd couting\\part_B_final\\train_data\\images’ # 上海图片数据集
SHHB_txt = ‘D:\\project\\crowd couting\\part_B_final\\train_data\\Part_B_train.txt’ # txt数据文件
pro_train = ‘D:\\project\\PyTorch-YOLOv3-master\\pro_train.txt’ # 把txt中的文件分为训练和验证集写到此文件中
pro_valid = ‘D:\\project\\PyTorch-YOLOv3-master\\pro_valid.txt’
# pro_train = ‘D:\\project\\zzzzzzz\\ag_image\\pro_train.txt’ # 把txt中的文件分为训练和验证集写到此文件中
# pro_valid = ‘D:\\project\\zzzzzzz\\ag_image\\pro_valid.txt’

# classes = [\'head\']# pre_define_categories = {}# for i, cls in enumerate(classes):#     pre_define_categories[cls] = i# only_care_pre_define_categories = True# only_care_pre_define_categories = Falseval_train(SHHB_img, SHHB_txt, pro_train, pro_valid)convert(pro_train, pro_valid)