划分数据集
将数据集划分为训练集和测试集
import os
import shutil
import random # 数据集路径
images_dir = './images'
targets_dir = './gt' # 划分后的路径
train_images_dir = './train/images'
test_images_dir = './test/images'
train_targets_dir = './train/gt'
test_targets_dir = './test/gt'os.makedirs(train_images_dir, exist_ok=True)
os.makedirs(test_images_dir, exist_ok=True)
os.makedirs(train_targets_dir, exist_ok=True)
os.makedirs(test_targets_dir, exist_ok=True) image_files = [f for f in os.listdir(images_dir) if os.path.isfile(os.path.join(images_dir, f))]
target_files = [f for f in os.listdir(targets_dir) if os.path.isfile(os.path.join(targets_dir, f))] train_ratio = 0.8
test_ratio = 1 - train_ratio combined = list(zip(image_files, target_files))
random.shuffle(combined)
image_files, target_files = zip(*combined) train_size = int(len(image_files) * train_ratio)
test_size = len(image_files) - train_size
train_images, test_images = image_files[:train_size], image_files[train_size:]
train_targets, test_targets = target_files[:train_size], target_files[train_size:] for img, tgt in zip(train_images, train_targets): shutil.copy(os.path.join(images_dir, img), os.path.join(train_images_dir, img)) shutil.copy(os.path.join(targets_dir, tgt), os.path.join(train_targets_dir, tgt)) for img, tgt in zip(test_images, test_targets): shutil.copy(os.path.join(images_dir, img), os.path.join(test_images_dir, img)) shutil.copy(os.path.join(targets_dir, tgt), os.path.join(test_targets_dir, tgt)) print("Training and test sets have been created successfully!")
视频成帧和数据标注
import imageiofilename = "./original.MP4"
vid = imageio.get_reader(filename, 'ffmpeg')try:for num, im in enumerate(vid):if (num / 50) and (num % 50) == 0: # 控制图像的输出张数；print(num // 50)imageio.imwrite('./jpg/{}.jpg'.format(num // 50), im)else:continue
except imageio.core.format.CannotReadFrameError or RuntimeError:pass
import cv2
import os"""
This code is used to:
1)对图片进行标注
2）生成对应的包含坐标信息的.txt文件
"""imgs_path = "./jpg/"
txt_path = "./label/"
files = os.listdir(imgs_path)
img = 0
coordinates = []def on_EVENT_LBUTTONDOWN(event, x, y, flags, param):if event == cv2.EVENT_LBUTTONDOWN:cv2.circle(img, (x, y), 4, (0, 255, 0), thickness=-1)coordinates.append([x, y])print([x,y])cv2.imshow("image", img)for file in files: # for i in range(80, len(files)):coordinates = []img = cv2.imread(imgs_path+file)cv2.namedWindow("image")cv2.setMouseCallback("image", on_EVENT_LBUTTONDOWN)cv2.imshow("image", img)cv2.waitKey(0)with open(txt_path+file.replace("jpg","txt"), "w+") as f:for coor in coordinates:f.write(str(coor[0])+" "+str(coor[1])+"\n") # 记录每个人头的坐标点# f.write(str(len(coordinates))) # 记录一张图像中的人头总数print(file+" is ok !"+"\n")
删除不匹配的文件
import osimg_folder = ".\\images"
txt_folder = ".\\gt"img_files = [f for f in os.listdir(img_folder) if f.endswith(".jpg")]
txt_files = [f for f in os.listdir(txt_folder) if f.endswith(".txt")]
img_file_names = [os.path.splitext(f)[0] for f in img_files]
for txt_file in txt_files:if os.path.splitext(txt_file)[0] not in img_file_names:os.remove(os.path.join(txt_folder, txt_file))print(f"Deleted {txt_file}")
：递归进阶练习)
)
