94 lines
2.9 KiB
Python
94 lines
2.9 KiB
Python
|
import os
|
||
|
|
|
||
|
|
# ----------------------------
|
||
|
|
# 配置
|
||
|
|
# ----------------------------
|
||
|
|
SOURCE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/resize_seg2" # 原始图片目录
|
||
|
|
TARGET_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/resize_seg3" # 输出目录
|
||
|
|
SUBSETS = ["train", "val", "test"]
|
||
|
|
|
||
|
|
# 全局 ROI (x, y, w, h)
|
||
|
|
GLOBAL_ROI = [562, 798, 1287, 453]
|
||
|
|
|
||
|
|
# ----------------------------
|
||
|
|
# 分割标签处理函数
|
||
|
|
# ----------------------------
|
||
|
|
def adjust_seg_labels(label_path, crop_coords, orig_size):
|
||
|
|
x1, y1, x2, y2 = crop_coords
|
||
|
|
crop_w = x2 - x1
|
||
|
|
crop_h = y2 - y1
|
||
|
|
new_labels = []
|
||
|
|
|
||
|
|
if not os.path.exists(label_path):
|
||
|
|
print(f"⚠️ 标签源文件不存在: {label_path}")
|
||
|
|
return new_labels
|
||
|
|
|
||
|
|
with open(label_path, 'r') as f:
|
||
|
|
lines = f.readlines()
|
||
|
|
|
||
|
|
orig_w, orig_h = orig_size
|
||
|
|
|
||
|
|
for line in lines:
|
||
|
|
parts = line.strip().split()
|
||
|
|
if len(parts) < 3 or len(parts) % 2 == 0:
|
||
|
|
continue # 至少类别 + 一个点(x,y)
|
||
|
|
|
||
|
|
cls = parts[0]
|
||
|
|
coords = list(map(float, parts[1:]))
|
||
|
|
|
||
|
|
new_coords = []
|
||
|
|
for i in range(0, len(coords), 2):
|
||
|
|
x = coords[i] * orig_w
|
||
|
|
y = coords[i+1] * orig_h
|
||
|
|
x -= x1
|
||
|
|
y -= y1
|
||
|
|
x_new = x / crop_w
|
||
|
|
y_new = y / crop_h
|
||
|
|
new_coords.extend([x_new, y_new])
|
||
|
|
|
||
|
|
new_line = cls + " " + " ".join([f"{c:.6f}" for c in new_coords])
|
||
|
|
new_labels.append(new_line)
|
||
|
|
|
||
|
|
return new_labels
|
||
|
|
|
||
|
|
# ----------------------------
|
||
|
|
# 主处理函数
|
||
|
|
# ----------------------------
|
||
|
|
def process_seg_labels():
|
||
|
|
x, y, w, h = GLOBAL_ROI
|
||
|
|
fixed_roi = (x, y, x + w, y + h)
|
||
|
|
|
||
|
|
for subset in SUBSETS:
|
||
|
|
src_dir = os.path.join(SOURCE_DIR, subset)
|
||
|
|
tgt_dir = os.path.join(TARGET_DIR, subset)
|
||
|
|
os.makedirs(tgt_dir, exist_ok=True)
|
||
|
|
|
||
|
|
for file in os.listdir(src_dir):
|
||
|
|
if not (file.endswith(".txt")):
|
||
|
|
continue
|
||
|
|
|
||
|
|
label_path = os.path.join(src_dir, file)
|
||
|
|
# 读取原图尺寸
|
||
|
|
img_name = os.path.splitext(file)[0] + ".jpg"
|
||
|
|
img_path = os.path.join(src_dir, img_name)
|
||
|
|
if not os.path.exists(img_path):
|
||
|
|
print(f"❌ 无法读取图片以获取尺寸: {img_path}")
|
||
|
|
continue
|
||
|
|
|
||
|
|
import cv2
|
||
|
|
img = cv2.imread(img_path)
|
||
|
|
if img is None:
|
||
|
|
print(f"❌ 无法读取图片: {img_path}")
|
||
|
|
continue
|
||
|
|
h_img, w_img = img.shape[:2]
|
||
|
|
|
||
|
|
new_labels = adjust_seg_labels(label_path, fixed_roi, (w_img, h_img))
|
||
|
|
tgt_label_path = os.path.join(tgt_dir, file)
|
||
|
|
with open(tgt_label_path, 'w') as f:
|
||
|
|
f.write("\n".join(new_labels))
|
||
|
|
|
||
|
|
print(f"✅ 标签处理完成: {subset}/{file}, 条数 {len(new_labels)}")
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
process_seg_labels()
|