194 lines
6.7 KiB
Python
194 lines
6.7 KiB
Python
|
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)
|