[CF]提交文件

src/pic.py

@@ -0,0 +1,88 @@
+import os
+import cv2
+def get_class_folder(catagetory, class_id, base_folder):
+    # 根据类别ID创建文件夹路径
+    class_folder = os.path.join(base_folder, catagetory[int(class_id)])
+    if not os.path.exists(class_folder):
+        os.makedirs(class_folder)
+    return class_folder
+
+def crop_and_save(image_path, label, class_folder):
+    # 读取图片
+    img = cv2.imread(image_path)
+    height, width = img.shape[:2]
+
+    # 解析标签
+    parts = label.strip().split()
+
+    x_center = float(parts[1]) * width
+    y_center = float(parts[2]) * height
+    img_width = float(parts[3]) * width
+    img_height = float(parts[4]) * height
+
+    # 计算左上角坐标
+    x1 = int(x_center - (img_width / 2))
+    y1 = int(y_center - (img_height / 2))
+
+    # 确保裁剪区域在图片范围内
+    x1 = max(0, x1)
+    y1 = max(0, y1)
+    x2 = min(width, x1 + img_width)
+    y2 = min(height, y1 + img_height)
+
+    # 裁剪图片
+    cropped_img = img[int(y1):int(y2), int(x1):int(x2)]
+
+    # 保存裁剪后的图片，文件名为数字递增
+    file_number = 1
+    file_path = os.path.join(class_folder, f"{file_number}.jpg")
+    while os.path.exists(file_path):
+        file_number += 1
+        file_path = os.path.join(class_folder, f"{file_number}.jpg")
+
+    cv2.imwrite(file_path, cropped_img)
+
+
+def main():
+    root = r"D:deskAdd_Building_20240616IR"
+    images_folder = os.path.join(root, 'images')
+    labels_folder = os.path.join(root, 'labels')
+    cropped_images_base_folder = os.path.join(root, 'cropped_images')
+    class_path = os.path.join(labels_folder, r"classes.txt")
+    category = {}
+    with open(class_path, 'r', encoding='utf-8') as f:
+        cls = f.readlines()
+        for i, c in enumerate(cls):
+            category[i] = c.strip()
+
+    # 确保裁剪图片的基文件夹存在
+    if not os.path.exists(cropped_images_base_folder):
+        os.makedirs(cropped_images_base_folder)
+
+    # 遍历images文件夹中的所有图片
+    for image_filename in os.listdir(images_folder):
+        if image_filename.lower().endswith(('.png', '.jpg', '.jpeg')):
+            # 构建图片和标签的完整路径
+            image_path = os.path.join(images_folder, image_filename)
+            label_path = os.path.join(labels_folder, os.path.splitext(image_filename)[0] + '.txt')
+
+            # 确保标签文件存在
+            if os.path.exists(label_path):
+                # 读取标签文件
+                with open(label_path, 'r') as file:
+                    labels = file.readlines()
+
+                # 遍历每个标签
+                for label in labels:
+                    # 获取类别ID
+                    class_id = int(label.strip().split()[0])
+
+                    # 获取类别文件夹
+                    class_folder = get_class_folder(category, class_id, cropped_images_base_folder)
+
+                    # 裁剪并保存图片
+                    crop_and_save(image_path, label, class_folder)
+
+
+if __name__ == '__main__':
+    main()

src/test.py

@@ -1,3 +1,3 @@
 from ultralytics import YOLO
 model = YOLO(r"C:\workspace\le-yolo\runs\detect\train26\weights\best.pt")
-results = model.predict("../res/3.mp4", show=True, save=True)
+results = model.predict("../res/5.mp4", show=True, save=True)