import os import cv2 import re import zipfile import shutil import numpy as np # ---------------------- 工具函数 ---------------------- def coord_to_num(x, y, cols, cell_w, cell_h): """根据像素坐标反算格子编号（从1开始）""" c = int(x // cell_w) r = int(y // cell_h) return r * cols + c + 1 def num_to_coord(num, cols, cell_w, cell_h): """网格编号转像素坐标""" n = num - 1 r, c = divmod(n, cols) x1, y1 = c * cell_w, r * cell_h x2, y2 = x1 + cell_w, y1 + cell_h return x1, y1, x2, y2 def convex_hull_poly(points): """计算点集凸包""" pts = np.array(points) hull = cv2.convexHull(pts) return hull.reshape(-1, 2).tolist() def draw_grid_on_image(image_path, grid_cells, cell_size=(108,102), save_path=None): """在图像上绘制网格编号和彩色框""" image = cv2.imread(image_path) if image is None: print(f"❌ 无法读取图片: {image_path}") return h, w = image.shape[:2] cell_w, cell_h = cell_size cols = w // cell_w overlay = image.copy() for cname, nums in grid_cells.items(): color = (np.random.randint(64,255), np.random.randint(64,255), np.random.randint(64,255)) for num in nums: x1, y1, x2, y2 = num_to_coord(num, cols, cell_w, cell_h) cv2.rectangle(overlay, (x1,y1), (x2,y2), color, -1) cv2.addWeighted(overlay, 0.4, image, 0.6, 0, image) # 绘制网格线 for i in range(0, w, cell_w): cv2.line(image, (i,0), (i,h), (100,100,100), 1) for j in range(0, h, cell_h): cv2.line(image, (0,j), (w,j), (100,100,100), 1) if save_path: cv2.imwrite(save_path, image) return image # ---------------------- YOLO-Seg → 网格编号 ---------------------- def yoloseg_to_grid_cells_fixed_v5(image_path, label_file, cell_size=(108,102), class_names=None): import numpy as np import cv2 img = cv2.imread(image_path) if img is None: raise ValueError(f"无法读取图片 {image_path}") h, w = img.shape[:2] cell_w, cell_h = cell_size cols = max(1, w // cell_w) class_cells = {} result_lines = [] with open(label_file, 'r', encoding='utf-8') as f: for line in f: parts = line.strip().split() if len(parts) < 5: continue cls_id = int(parts[0]) coords = [float(x) for x in parts[1:]] if len(coords) % 2 != 0: coords = coords[:-1] if len(coords) < 6: continue poly = np.array(coords, dtype=np.float32).reshape(-1, 2) poly[:, 0] *= w poly[:, 1] *= h x_min, y_min = poly.min(axis=0) x_max, y_max = poly.max(axis=0) col_start = int(x_min // cell_w) col_end = int(x_max // cell_w) row_start = int(y_min // cell_h) row_end = int(y_max // cell_h) # 直接生成网格编号 cells = [r*cols + c + 1 for r in range(row_start, row_end+1) for c in range(col_start, col_end+1)] if class_names and isinstance(class_names, list): cname = class_names[cls_id] if cls_id < len(class_names) else str(cls_id) else: cname = str(cls_id) if cname not in class_cells: class_cells[cname] = set() class_cells[cname].update(cells) # 输出 _grid.txt for cname, nums in class_cells.items(): sorted_nums = sorted(nums) ids_str = '-'.join(map(str, sorted_nums)) + '-' result_lines.append(f"{cname} {ids_str}") return '\n'.join(result_lines), class_cells # ---------------------- 主函数 ---------------------- def process_zip_yoloseg_with_draw(zip_path, output_dir, cell_size=(108,102), class_names=None, make_zip=True): """ 处理 YOLO-Seg 数据集压缩包 → 网格编号 + 可视化 支持标签在不同子文件夹 """ if not os.path.exists(zip_path): raise FileNotFoundError(f"{zip_path} 不存在") os.makedirs(output_dir, exist_ok=True) # 解压 with zipfile.ZipFile(zip_path, 'r') as zip_ref: zip_ref.extractall(output_dir) # 枚举所有图片 for root, _, files in os.walk(output_dir): for file in files: if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')): image_path = os.path.join(root, file) base_name = os.path.splitext(file)[0] # 在所有子目录中寻找对应标签 label_file = None for subroot, _, subfiles in os.walk(output_dir): for sf in subfiles: if sf == base_name + ".txt": label_file = os.path.join(subroot, sf) break if label_file: break if not label_file: print(f"⚠️ 找不到标签文件: {base_name}.txt") continue try: output_str, class_cells = yoloseg_to_grid_cells_fixed_v5( image_path, label_file, cell_size=cell_size, class_names=class_names ) except Exception as e: print(f"❌ 转换失败 {file}: {e}") continue # 写入 _grid.txt out_txt_file = os.path.join(root, base_name + "_grid.txt") with open(out_txt_file, 'w', encoding='utf-8') as f: f.write(output_str) # 绘制结果图 out_img_file = os.path.join(root, base_name + "_grid.jpg") draw_grid_on_image(image_path, class_cells, cell_size=cell_size, save_path=out_img_file) # 打包结果 if make_zip: zip_out_path = os.path.splitext(output_dir.rstrip("/\\"))[0] + "_processed.zip" shutil.make_archive(os.path.splitext(zip_out_path)[0], 'zip', output_dir) print(f"✅ 处理完成，已生成压缩包: {zip_out_path}") return zip_out_path else: print(f"✅ 处理完成，输出目录: {output_dir}") return output_dir # ---------------------- 示例调用 ---------------------- if __name__ == "__main__": # zip_path = r"D:\work\develop\LF-where\1760968917215-1760968915394853977_all_zip.zip" # zip_path = r"C:\Users\14867\Downloads\1761100847201-1761100834344385200_all_zip.zip" zip_path = r"/test/predictions-20250711-111516-531.zip" output_dir = "../test/out" classes = ['裂缝','横向裂缝','纵向裂缝',"修补","坑洞"] process_zip_yoloseg_with_draw(zip_path, output_dir, class_names=classes)