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)
|