82 lines
3.0 KiB
Python
82 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
import os, cv2, numpy as np
|
|
|
|
# ========== 配置 ==========
|
|
roi_label_dir = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/label" # 640x640 ROI 标签
|
|
original_image_dir = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/class1"
|
|
roi_coords_file = "/home/hx/yolo/zhuangtai_class_cls/roi_coordinates/1_rois.txt" # 每行: x,y,w,h
|
|
output_label_dir = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/classdata3_orig_norm_labels"
|
|
debug_vis = True
|
|
debug_vis_dir = output_label_dir + "_vis"
|
|
# ==========================
|
|
|
|
os.makedirs(output_label_dir, exist_ok=True)
|
|
if debug_vis:
|
|
os.makedirs(debug_vis_dir, exist_ok=True)
|
|
|
|
# ----------- 加载 ROI -----------
|
|
with open(roi_coords_file, 'r') as f:
|
|
roi_x, roi_y, roi_w, roi_h = map(int, f.readline().strip().split(','))
|
|
|
|
# ----------- 处理标签 -----------
|
|
for label_file in os.listdir(roi_label_dir):
|
|
if not label_file.lower().endswith('.txt'):
|
|
continue
|
|
base_name = os.path.splitext(label_file)[0]
|
|
|
|
# 对应原图
|
|
orig_img_path = None
|
|
for ext in ('.jpg','.jpeg','.png','.bmp'):
|
|
tmp_path = os.path.join(original_image_dir, base_name + ext)
|
|
if os.path.exists(tmp_path):
|
|
orig_img_path = tmp_path
|
|
break
|
|
if orig_img_path is None:
|
|
print(f"[SKIP] 未找到原图: {base_name}")
|
|
continue
|
|
|
|
img = cv2.imread(orig_img_path)
|
|
if img is None:
|
|
print(f"[SKIP] 无法读取原图: {orig_img_path}")
|
|
continue
|
|
h_img, w_img = img.shape[:2]
|
|
|
|
# 读取 ROI 标签
|
|
in_path = os.path.join(roi_label_dir, label_file)
|
|
lines = [ln.strip() for ln in open(in_path, 'r') if ln.strip()]
|
|
out_lines = []
|
|
|
|
for ln in lines:
|
|
parts = ln.split()
|
|
cls = parts[0]
|
|
coords = list(map(float, parts[1:]))
|
|
mapped = []
|
|
for i in range(0, len(coords), 2):
|
|
x_n, y_n = coords[i], coords[i+1] # ROI 0~1
|
|
x_abs = roi_x + x_n * roi_w
|
|
y_abs = roi_y + y_n * roi_h
|
|
x_abs = min(max(x_abs, 0.0), w_img-1)
|
|
y_abs = min(max(y_abs, 0.0), h_img-1)
|
|
mapped.append((x_abs / w_img, y_abs / h_img))
|
|
out_line = cls + " " + " ".join(f"{x:.6f} {y:.6f}" for x, y in mapped)
|
|
out_lines.append(out_line)
|
|
|
|
out_path = os.path.join(output_label_dir, label_file)
|
|
with open(out_path, 'w') as fw:
|
|
fw.write("\n".join(out_lines))
|
|
print(f"[OK] {out_path}")
|
|
|
|
# 可视化
|
|
if debug_vis and out_lines:
|
|
vis = img.copy()
|
|
cv2.rectangle(vis, (roi_x, roi_y), (roi_x+roi_w-1, roi_y+roi_h-1), (0,165,255), 2)
|
|
for ol in out_lines:
|
|
pts = np.array([[int(float(parts[i])*w_img), int(float(parts[i+1])*h_img)]
|
|
for parts in [ol.split()] for i in range(1, len(parts), 2)], np.int32)
|
|
if len(pts) >= 3:
|
|
cv2.polylines(vis, [pts], True, (0,0,255), 2)
|
|
vis_name = os.path.join(debug_vis_dir, os.path.basename(orig_img_path))
|
|
cv2.imwrite(vis_name, vis)
|
|
|
|
print("🎉 ROI 标签已还原到原图归一化完成！")
|