250 lines
8.9 KiB
Python
250 lines
8.9 KiB
Python
|
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()
|
|||
|
|
|
|||
|
|
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):
|
|||
|
|
"""
|
|||
|
|
强力匹配：
|
|||
|
|
- 去掉 _PartClass
|
|||
|
|
- 去掉 .txt
|
|||
|
|
- 如果有 .jpg 在 TXT 名里，也去掉
|
|||
|
|
- 模糊匹配核心名和图片名
|
|||
|
|
"""
|
|||
|
|
txt_name = os.path.basename(txt_path).lower()
|
|||
|
|
# 去掉 _partclass 和 .txt
|
|||
|
|
base_name = re.sub(r'(_partclass)?\.txt$', '', txt_name)
|
|||
|
|
# 再去掉可能残留的 .jpg
|
|||
|
|
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)
|
|||
|
|
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]
|
|||
|
|
os.makedirs(os.path.join(output_root,"images"), exist_ok=True)
|
|||
|
|
os.makedirs(os.path.join(output_root,"labels"), exist_ok=True)
|
|||
|
|
os.makedirs(os.path.join(output_root,"visual"), 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)
|
|||
|
|
print(f"✅ 已处理像素点 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)
|
|||
|
|
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]
|
|||
|
|
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)
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
# ------------------ 批量处理 ------------------
|
|||
|
|
def batch_process_txt_first(input_root, class_map, output_root="output", max_workers=4):
|
|||
|
|
os.makedirs(os.path.join(output_root,"images"), exist_ok=True)
|
|||
|
|
os.makedirs(os.path.join(output_root,"labels"), exist_ok=True)
|
|||
|
|
os.makedirs(os.path.join(output_root,"visual"), exist_ok=True)
|
|||
|
|
os.makedirs(os.path.join(output_root,"highlighted"), exist_ok=True)
|
|||
|
|
|
|||
|
|
# 收集所有 TXT 文件
|
|||
|
|
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))
|
|||
|
|
|
|||
|
|
success_count, fail_count = 0, 0
|
|||
|
|
log_lines = []
|
|||
|
|
fail_logs = []
|
|||
|
|
|
|||
|
|
def process_single(txt_path):
|
|||
|
|
nonlocal success_count, fail_count
|
|||
|
|
img_path = find_matching_image(txt_path, input_root)
|
|||
|
|
if img_path:
|
|||
|
|
try:
|
|||
|
|
if "_partclass" in txt_path.lower():
|
|||
|
|
status = process_grid_txt(img_path, txt_path, class_map, output_root)
|
|||
|
|
log_lines.append(f"{os.path.basename(txt_path)} -> Grid TXT processed with {os.path.basename(img_path)}")
|
|||
|
|
else:
|
|||
|
|
status = process_pixel_txt(img_path, txt_path, class_map, output_root)
|
|||
|
|
log_lines.append(f"{os.path.basename(txt_path)} -> Pixel TXT processed with {os.path.basename(img_path)}")
|
|||
|
|
if status:
|
|||
|
|
success_count += 1
|
|||
|
|
else:
|
|||
|
|
fail_count += 1
|
|||
|
|
fail_logs.append(f"{os.path.basename(txt_path)} -> Processed but no valid labels generated")
|
|||
|
|
except Exception as e:
|
|||
|
|
fail_count += 1
|
|||
|
|
fail_logs.append(f"{os.path.basename(txt_path)} -> Processing error: {e}")
|
|||
|
|
else:
|
|||
|
|
fail_count += 1
|
|||
|
|
fail_logs.append(f"{os.path.basename(txt_path)} -> No matching image found")
|
|||
|
|
|
|||
|
|
from concurrent.futures import ThreadPoolExecutor
|
|||
|
|
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|||
|
|
executor.map(process_single, txt_files)
|
|||
|
|
|
|||
|
|
# 写入日志
|
|||
|
|
log_file = os.path.join(output_root, "process_log.txt")
|
|||
|
|
with open(log_file, "w", encoding="utf-8") as f:
|
|||
|
|
f.write("\n".join(log_lines + ["\n失败文件:"] + fail_logs))
|
|||
|
|
|
|||
|
|
print(f"\n✅ 批量处理完成: 成功 {success_count}, 失败 {fail_count}")
|
|||
|
|
if fail_logs:
|
|||
|
|
print("⚠️ 失败文件及原因如下:")
|
|||
|
|
for line in fail_logs:
|
|||
|
|
print(line)
|
|||
|
|
print(f"📄 处理日志已保存: {log_file}")
|
|||
|
|
|
|||
|
|
# ------------------ 主程序 ------------------
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
input_root = r"D:\work\develop\LF-where\01"
|
|||
|
|
output_root = r"D:\work\develop\LF-where\out"
|
|||
|
|
|
|||
|
|
class_map = {
|
|||
|
|
"裂缝": 0,
|
|||
|
|
"横向裂缝": 1,
|
|||
|
|
"纵向裂缝": 2,
|
|||
|
|
"修补": 3,
|
|||
|
|
"坑洞": 4,
|
|||
|
|
"网裂": 5,
|
|||
|
|
"破碎板":6,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
batch_process_txt_first(input_root, class_map, output_root, max_workers=8)
|