import os import re import cv2 import numpy as np import shutil from concurrent.futures import ThreadPoolExecutor # ------------------ 工具函数 ------------------ def clean_filename(name): name = name.strip() name = re.sub(r'[\s\r\n\t]+', '', name) return name.lower() def num_to_coord(num, cols, cell_width, cell_height, offset=1): n = num - 1 + offset r = n // cols c = n % cols x1 = c * cell_width y1 = r * cell_height x2 = x1 + cell_width y2 = y1 + cell_height return x1, y1, x2, y2 def polygon_to_yolo(poly, img_width, img_height): flat = [coord for point in poly for coord in point] return [flat[i] / (img_width if i % 2 == 0 else img_height) for i in range(len(flat))] def convex_hull_poly(points): if not points: return [] pts = np.array(points, dtype=np.int32) hull = cv2.convexHull(pts) return hull.reshape(-1, 2).tolist() def expand_polygon(poly, expand_px=3): """多边形顶点膨胀 expand_px 像素""" if not poly: return [] pts = np.array(poly, dtype=np.int32) cx = np.mean(pts[:, 0]) cy = np.mean(pts[:, 1]) vec = pts - np.array([[cx, cy]]) norm = np.linalg.norm(vec, axis=1, keepdims=True) norm[norm == 0] = 1 vec_unit = vec / norm pts_expanded = pts + (vec_unit * expand_px).astype(int) return pts_expanded.tolist() color_map = { 0: (0, 255, 255), 1: (255, 0, 255), 2: (0, 255, 0), 3: (255, 0, 0), 4: (0, 0, 255), 5: (255, 255, 0), 6: (128, 128, 0), 7: (128, 0, 128), 8: (0, 128, 128), 9: (128, 128, 128), 10: (0, 0, 128), 11: (0, 128, 0) } # ------------------ 匹配图片 ------------------ def find_matching_image(txt_path, input_root): txt_name = os.path.basename(txt_path).lower() base_name = re.sub(r'(_partclass)?\.txt$', '', txt_name) base_name = re.sub(r'\.jpg$', '', base_name) for root, _, files in os.walk(input_root): for f in files: if f.lower().endswith((".jpg", ".jpeg", ".png")): img_base = os.path.splitext(f)[0].lower() if base_name == img_base: return os.path.join(root, f) return None # ------------------ 原有处理函数（带多边形膨胀） ------------------ def process_pixel_txt(img_path, txt_path, class_map, output_root): image = cv2.imread(img_path) if image is None: return False h, w = image.shape[:2] vis_img = image.copy() yolo_labels = [] unknown_labels = set() with open(txt_path, "r", encoding="utf-8") as f: for line in f: parts = line.strip().split() if len(parts) < 5: continue try: x, y, w_box, h_box = map(int, parts[:4]) except: continue label = parts[4] cls_id = class_map.get(label, -1) if cls_id == -1: unknown_labels.add(label) continue poly = [(x, y), (x + w_box, y), (x + w_box, y + h_box), (x, y + h_box)] hull = convex_hull_poly(poly) hull = expand_polygon(hull, expand_px=3) yolo_labels.append(f"{cls_id} " + " ".join(map(str, polygon_to_yolo(hull, w, h)))) cv2.polylines(vis_img, [np.array(hull, np.int32)], True, color=color_map.get(cls_id, (255, 255, 255)), thickness=2) if unknown_labels: print(f"⚠️ 未知类别 {unknown_labels} 在文件: {txt_path}") if not yolo_labels: return False base = os.path.splitext(os.path.basename(img_path))[0] folder_img = os.path.join(output_root, "images") os.makedirs(folder_img, exist_ok=True) shutil.copy2(img_path, os.path.join(folder_img, os.path.basename(img_path))) folder_labels = os.path.join(output_root, "labels") os.makedirs(folder_labels, exist_ok=True) with open(os.path.join(folder_labels, base + ".txt"), "w", encoding="utf-8") as f: f.write("\n".join(yolo_labels)) folder_vis = os.path.join(output_root, "visual") os.makedirs(folder_vis, exist_ok=True) cv2.imwrite(os.path.join(folder_vis, base + "-visual.jpg"), vis_img) print(f"✅ 已处理 Pixel TXT: {base}") return True def process_grid_txt(img_path, txt_path, class_map, output_root): image = cv2.imread(img_path) if image is None: return False h, w = image.shape[:2] cell_width, cell_height = 108, 102 cols = max(1, w // cell_width) vis_img = image.copy() overlay = image.copy() alpha = 0.5 yolo_labels = [] with open(txt_path, "r", encoding="utf-8") as f: for line in f: line = line.strip() if not line: continue numbers = re.findall(r"(\d+)(?=-|$)", line.split()[-1]) numbers = [int(n) for n in numbers] cname = None for key in class_map.keys(): if line.startswith(key): cname = key break if cname is None or not numbers: continue for num in numbers: x1, y1, x2, y2 = num_to_coord(num, cols, cell_width, cell_height) cv2.rectangle(overlay, (x1, y1), (x2, y2), color_map.get(class_map[cname], (128, 128, 128)), -1) cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image) points = [] for num in numbers: x1, y1, x2, y2 = num_to_coord(num, cols, cell_width, cell_height) points.extend([(x1, y1), (x2, y1), (x2, y2), (x1, y2)]) hull = convex_hull_poly(points) hull = expand_polygon(hull, expand_px=3) cls_id = class_map[cname] pts = np.array(hull, np.int32).reshape((-1, 1, 2)) cv2.polylines(vis_img, [pts], True, color_map.get(cls_id, (128, 128, 128)), 2) yolo_labels.append(f"{cls_id} " + " ".join(map(str, polygon_to_yolo(hull, w, h)))) if not yolo_labels: return False base = os.path.splitext(os.path.basename(img_path))[0] for folder_name in ["images", "labels", "visual", "highlighted"]: folder = os.path.join(output_root, folder_name) os.makedirs(folder, exist_ok=True) shutil.copy2(img_path, os.path.join(output_root, "images", os.path.basename(img_path))) with open(os.path.join(output_root, "labels", base + ".txt"), "w", encoding="utf-8") as f: f.write("\n".join(yolo_labels)) cv2.imwrite(os.path.join(output_root, "visual", base + "-visual.jpg"), vis_img) cv2.imwrite(os.path.join(output_root, "highlighted", base + "-highlighted.jpg"), image) print(f"✅ 已处理 Grid TXT: {base}") return True # ------------------ 批量处理函数 ------------------ def batch_process_txt_first(input_root, output_root, mode_type="5211", max_workers=4): if mode_type == "5211": class_maps = { "asphalt": class_map_asphalt, "cream": class_map_cream, "gravel": class_map_gravel, } elif mode_type == "5210": class_maps = { "asphalt": class_map_asphalt_road, "cream": class_map_cream_road, "gravel": class_map_gravel, } else: raise ValueError("mode_type 必须是 '5210' 或 '5211'") txt_files = [] for root, _, files in os.walk(input_root): for file in files: if file.lower().endswith(".txt"): txt_files.append(os.path.join(root, file)) def process_single(txt_path): img_path = find_matching_image(txt_path, input_root) if not img_path: print(f"⚠️ 未找到匹配图片: {txt_path}") return with open(txt_path, "r", encoding="utf-8") as f: content = f.read().lower() if any(k in content for k in ["水泥", "cement", "cream"]): road_type = "cream" elif any(k in content for k in ["沥青", "asphalt"]): road_type = "asphalt" else: road_type = "gravel" selected_map = class_maps[road_type] output_folder = os.path.join(output_root, road_type) os.makedirs(output_folder, exist_ok=True) if "_partclass" in txt_path.lower(): process_grid_txt(img_path, txt_path, selected_map, output_folder) else: process_pixel_txt(img_path, txt_path, selected_map, output_folder) with ThreadPoolExecutor(max_workers=max_workers) as executor: executor.map(process_single, txt_files) # ------------------ 类别映射 ------------------ class_map_asphalt = { # "纵向裂缝": 0, "横向裂缝": 1, "网裂": 2, "坑槽": 3, "松散": 4, "龟裂": 0, "块状裂缝": 1, "纵向裂缝": 2, "横向裂缝": 3, "沉陷": 4, "车辙": 5, "波浪拥包": 6, "坑槽": 7, "松散": 8, "泛油": 9, "修补": 10, } class_map_cream = { "破碎板": 0, "裂缝": 1, "坑洞": 2, "露骨": 3, "错台": 4, "拱起": 5, } class_map_gravel = { "坑槽": 0, "沉陷": 1, "车辙": 2, "波浪搓板": 3, } class_map_asphalt_road = { "龟裂": 0, "块状裂缝": 1, "纵向裂缝": 2, "横向裂缝": 3, "沉陷": 4, "车辙": 5, "波浪拥包": 6, "坑槽": 7, "松散": 8, "泛油": 9, "修补": 10, } class_map_cream_road = { "破碎板": 0, "裂缝": 1, "板角断裂": 2, "错台": 3, "拱起": 4, "边角剥落": 5, "接缝料损坏": 6, "坑洞": 7, "唧泥": 8, "露骨": 9, "修补": 10, } # ------------------ 主程序 ------------------ if __name__ == "__main__": input_root = r"D:\work\develop\LF-where\01" output_root = r"D:\work\develop\LF-where\out" mode_type = "5211" # 5211 或 5210 batch_process_txt_first(input_root, output_root, mode_type, max_workers=8)