bushu
This commit is contained in:
BIN
zhuangtai_class_cls/best.pt
Normal file
BIN
zhuangtai_class_cls/best.pt
Normal file
Binary file not shown.
@ -75,7 +75,7 @@ def select_roi(image_path):
|
||||
if key == ord('s'):
|
||||
# 保存坐标
|
||||
base_name = os.path.splitext(os.path.basename(image_path))[0]
|
||||
save_path = os.path.join(save_dir, f"{base_name}_rois.txt") # 修改了扩展名为 .txt
|
||||
save_path = os.path.join(save_dir, f"{base_name}_rois1.txt") # 修改了扩展名为 .txt
|
||||
save_rois_to_txt(roi_list, save_path) # 使用新的保存函数
|
||||
|
||||
elif key == ord('n'):
|
||||
@ -1,84 +0,0 @@
|
||||
import cv2
|
||||
import os
|
||||
|
||||
# 配置路径
|
||||
original_images_parent_dir = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/classdata3/val" # 含 class0, class1 等
|
||||
roi_coords_file = "./roi_coordinates/1_rois.txt" # 你手动标注的唯一一个 ROI 文件
|
||||
output_parent_dir = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/classdata3/classdata3_cropped"
|
||||
target_size = 640
|
||||
|
||||
os.makedirs(output_parent_dir, exist_ok=True)
|
||||
|
||||
def load_global_rois(txt_path):
|
||||
"""加载全局 ROI 坐标"""
|
||||
rois = []
|
||||
if not os.path.exists(txt_path):
|
||||
print(f"❌ ROI 文件不存在: {txt_path}")
|
||||
return rois
|
||||
with open(txt_path, 'r') as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line:
|
||||
try:
|
||||
x, y, w, h = map(int, line.split(','))
|
||||
rois.append((x, y, w, h))
|
||||
print(f"📌 加载 ROI: (x={x}, y={y}, w={w}, h={h})")
|
||||
except Exception as e:
|
||||
print(f"⚠️ 无法解析 ROI 行: {line}, 错误: {e}")
|
||||
return rois
|
||||
|
||||
# 加载全局 ROI
|
||||
rois = load_global_rois(roi_coords_file)
|
||||
if len(rois) == 0:
|
||||
print("❌ 没有加载到任何有效的 ROI 坐标,程序退出。")
|
||||
exit(1)
|
||||
|
||||
# 遍历所有子目录和图像
|
||||
for class_name in os.listdir(original_images_parent_dir):
|
||||
class_dir = os.path.join(original_images_parent_dir, class_name)
|
||||
if not os.path.isdir(class_dir):
|
||||
continue
|
||||
|
||||
# 创建输出类别目录
|
||||
output_class_dir = os.path.join(output_parent_dir, class_name)
|
||||
os.makedirs(output_class_dir, exist_ok=True)
|
||||
|
||||
print(f"🔄 处理类别: {class_name}")
|
||||
|
||||
for img_file in os.listdir(class_dir):
|
||||
if not img_file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
|
||||
continue
|
||||
|
||||
img_path = os.path.join(class_dir, img_file)
|
||||
base_name, ext = os.path.splitext(img_file)
|
||||
|
||||
# 读取图像
|
||||
img = cv2.imread(img_path)
|
||||
if img is None:
|
||||
print(f"❌ 无法读取图像: {img_path}")
|
||||
continue
|
||||
|
||||
# 对每一个 ROI 进行裁剪
|
||||
for i, (x, y, w, h) in enumerate(rois):
|
||||
# 检查坐标是否越界
|
||||
h_img, w_img = img.shape[:2]
|
||||
if x < 0 or y < 0 or x + w > w_img or y + h > h_img:
|
||||
print(f"⚠️ ROI 越界,跳过: {x},{y},{w},{h} in {img_path}")
|
||||
continue
|
||||
|
||||
roi_img = img[y:y+h, x:x+w]
|
||||
if roi_img.size == 0:
|
||||
print(f"⚠️ 空 ROI 区域: {x},{y},{w},{h} in {img_path}")
|
||||
continue
|
||||
|
||||
# resize
|
||||
roi_resized = cv2.resize(roi_img, (target_size, target_size), interpolation=cv2.INTER_AREA)
|
||||
|
||||
# 保存文件名:原名 + _roi0, _roi1...
|
||||
suffix = f"_roi{i}" if len(rois) > 1 else ""
|
||||
save_filename = f"{base_name}{suffix}{ext}"
|
||||
save_path = os.path.join(output_class_dir, save_filename)
|
||||
cv2.imwrite(save_path, roi_resized)
|
||||
print(f"✅ 保存: {save_path}")
|
||||
|
||||
print("🎉 所有图像已根据全局 ROI 裁剪并保存完成!")
|
||||
58
zhuangtai_class_cls/resize_dataset_image.py
Normal file
58
zhuangtai_class_cls/resize_dataset_image.py
Normal file
@ -0,0 +1,58 @@
|
||||
import os
|
||||
import cv2
|
||||
|
||||
# ----------------------------
|
||||
# 配置
|
||||
# ----------------------------
|
||||
SOURCE_ROOT_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/classdata2" # 原始图片根目录
|
||||
TARGET_ROOT_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/classdata4" # 输出根目录
|
||||
CLASSES = ["class0", "class1", "class2", "class3", "class4"] # 类别列表
|
||||
TARGET_SIZE = 640 # resize 尺寸
|
||||
SUBSETS = ["train", "val", "test"]
|
||||
|
||||
# ----------------------------
|
||||
# 全局 ROI (x, y, w, h)
|
||||
# ----------------------------
|
||||
GLOBAL_ROI = [859,810,696,328]
|
||||
|
||||
# ----------------------------
|
||||
# 主处理函数
|
||||
# ----------------------------
|
||||
def process_images():
|
||||
x, y, w, h = GLOBAL_ROI
|
||||
for subset in SUBSETS:
|
||||
for class_dir in CLASSES:
|
||||
src_dir = os.path.join(SOURCE_ROOT_DIR, subset, class_dir)
|
||||
tgt_dir = os.path.join(TARGET_ROOT_DIR, subset, class_dir)
|
||||
os.makedirs(tgt_dir, exist_ok=True)
|
||||
|
||||
if not os.path.exists(src_dir):
|
||||
print(f"警告: 源目录 {src_dir} 不存在,跳过")
|
||||
continue
|
||||
|
||||
for file in os.listdir(src_dir):
|
||||
if not (file.endswith(".jpg") or file.endswith(".png")):
|
||||
continue
|
||||
|
||||
img_path = os.path.join(src_dir, file)
|
||||
img = cv2.imread(img_path)
|
||||
if img is None:
|
||||
print(f"❌ 无法读取图片: {img_path}")
|
||||
continue
|
||||
|
||||
h_img, w_img = img.shape[:2]
|
||||
x1, y1 = max(0, x), max(0, y)
|
||||
x2, y2 = min(w_img, x + w), min(h_img, y + h)
|
||||
|
||||
cropped = img[y1:y2, x1:x2]
|
||||
if cropped.size == 0:
|
||||
print(f"❌ 裁剪结果为空: {file}")
|
||||
continue
|
||||
|
||||
resized = cv2.resize(cropped, (TARGET_SIZE, TARGET_SIZE))
|
||||
tgt_path = os.path.join(tgt_dir, file)
|
||||
cv2.imwrite(tgt_path, resized)
|
||||
print(f"✅ 图片处理完成: {subset}/{class_dir}/{file}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
process_images()
|
||||
@ -97,9 +97,9 @@ def classify_and_save_images(model_path, input_folder, output_root, roi_file, ta
|
||||
# 主程序
|
||||
# ---------------------------
|
||||
if __name__ == "__main__":
|
||||
model_path = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/cls_resize/exp_cls2/weights/best.pt"
|
||||
model_path = r"best.pt"
|
||||
input_folder = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/f6"
|
||||
output_root = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/class111"
|
||||
output_root = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/class44"
|
||||
roi_file = "./roi_coordinates/1_rois.txt" # 训练时使用的 ROI 文件
|
||||
target_size = 640
|
||||
|
||||
|
||||
@ -134,8 +134,8 @@ def batch_classify_images(model_path, input_folder, output_root, roi_file, targe
|
||||
# ---------------------------
|
||||
if __name__ == "__main__":
|
||||
model_path = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/cls_resize/exp_cls2/weights/best.pt"
|
||||
input_folder = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/f6"
|
||||
output_root = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/classified"
|
||||
input_folder = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1000"
|
||||
output_root = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1000/classified"
|
||||
roi_file = "./roi_coordinates/1_rois.txt"
|
||||
target_size = 640
|
||||
threshold = 0.4 # 可调节的比例系数
|
||||
168
zhuangtai_class_cls/resize_tuili_image_main.py
Normal file
168
zhuangtai_class_cls/resize_tuili_image_main.py
Normal file
@ -0,0 +1,168 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
import cv2
|
||||
import numpy as np
|
||||
from ultralytics import YOLO
|
||||
|
||||
# ---------------------------
|
||||
# 类别映射
|
||||
# ---------------------------
|
||||
CLASS_NAMES = {
|
||||
0: "未堆料",
|
||||
1: "小堆料",
|
||||
2: "大堆料",
|
||||
3: "未浇筑满",
|
||||
4: "浇筑满"
|
||||
}
|
||||
|
||||
# ---------------------------
|
||||
# 加载 ROI 列表
|
||||
# ---------------------------
|
||||
def load_global_rois(txt_path):
|
||||
rois = []
|
||||
if not os.path.exists(txt_path):
|
||||
print(f"❌ ROI 文件不存在: {txt_path}")
|
||||
return rois
|
||||
with open(txt_path, 'r') as f:
|
||||
for line in f:
|
||||
s = line.strip()
|
||||
if s:
|
||||
try:
|
||||
x, y, w, h = map(int, s.split(','))
|
||||
rois.append((x, y, w, h))
|
||||
except Exception as e:
|
||||
print(f"无法解析 ROI 行 '{s}': {e}")
|
||||
return rois
|
||||
|
||||
# ---------------------------
|
||||
# 裁剪并 resize ROI
|
||||
# ---------------------------
|
||||
def crop_and_resize(img, rois, target_size=640):
|
||||
crops = []
|
||||
h_img, w_img = img.shape[:2]
|
||||
for i, (x, y, w, h) in enumerate(rois):
|
||||
if x < 0 or y < 0 or x + w > w_img or y + h > h_img:
|
||||
continue
|
||||
roi = img[y:y+h, x:x+w]
|
||||
roi_resized = cv2.resize(roi, (target_size, target_size), interpolation=cv2.INTER_AREA)
|
||||
crops.append((roi_resized, i))
|
||||
return crops
|
||||
|
||||
# ---------------------------
|
||||
# class1/class2 加权判断
|
||||
# ---------------------------
|
||||
def weighted_small_large(pred_probs, threshold=0.4, w1=0.3, w2=0.7):
|
||||
p1 = float(pred_probs[1])
|
||||
p2 = float(pred_probs[2])
|
||||
total = p1 + p2
|
||||
if total > 0:
|
||||
score = (w1 * p1 + w2 * p2) / total
|
||||
else:
|
||||
score = 0.0
|
||||
final_class = "大堆料" if score >= threshold else "小堆料"
|
||||
return final_class, score, p1, p2
|
||||
|
||||
# ---------------------------
|
||||
# 单张图片推理函数
|
||||
# ---------------------------
|
||||
def classify_image_weighted(image, model, threshold=0.4):
|
||||
results = model(image)
|
||||
pred_probs = results[0].probs.data.cpu().numpy().flatten()
|
||||
class_id = int(pred_probs.argmax())
|
||||
confidence = float(pred_probs[class_id])
|
||||
class_name = CLASS_NAMES.get(class_id, f"未知类别({class_id})")
|
||||
|
||||
# class1/class2 使用加权得分
|
||||
if class_id in [1, 2]:
|
||||
final_class, score, p1, p2 = weighted_small_large(pred_probs, threshold=threshold)
|
||||
else:
|
||||
final_class = class_name
|
||||
score = confidence
|
||||
p1 = float(pred_probs[1])
|
||||
p2 = float(pred_probs[2])
|
||||
|
||||
return final_class, score, p1, p2
|
||||
|
||||
# ---------------------------
|
||||
# 批量推理主函数
|
||||
# ---------------------------
|
||||
def batch_classify_images(model_path, input_folder, output_root, roi_file, target_size=640, threshold=0.5):
|
||||
# 加载模型
|
||||
model = YOLO(model_path)
|
||||
|
||||
# 确保输出根目录存在
|
||||
output_root = Path(output_root)
|
||||
output_root.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 为所有类别创建目录
|
||||
class_dirs = {}
|
||||
for name in CLASS_NAMES.values():
|
||||
d = output_root / name
|
||||
d.mkdir(exist_ok=True)
|
||||
class_dirs[name] = d
|
||||
|
||||
rois = load_global_rois(roi_file)
|
||||
if not rois:
|
||||
print("❌ 没有有效 ROI,退出")
|
||||
return
|
||||
|
||||
# 遍历图片
|
||||
for img_path in Path(input_folder).glob("*.*"):
|
||||
if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.tif']:
|
||||
continue
|
||||
try:
|
||||
img = cv2.imread(str(img_path))
|
||||
if img is None:
|
||||
continue
|
||||
|
||||
crops = crop_and_resize(img, rois, target_size)
|
||||
|
||||
for roi_resized, roi_idx in crops:
|
||||
final_class, score, p1, p2 = classify_image_weighted(roi_resized, model, threshold=threshold)
|
||||
|
||||
# 文件名中保存 ROI、类别、加权分数、class1/class2 置信度
|
||||
suffix = f"_roi{roi_idx}_{final_class}_score{score:.2f}_p1{p1:.2f}_p2{p2:.2f}"
|
||||
dst_path = class_dirs[final_class] / f"{img_path.stem}{suffix}{img_path.suffix}"
|
||||
cv2.imwrite(dst_path, roi_resized)
|
||||
print(f"{img_path.name}{suffix} -> {final_class} (score={score:.2f}, p1={p1:.2f}, p2={p2:.2f})")
|
||||
|
||||
except Exception as e:
|
||||
print(f"处理失败 {img_path.name}: {e}")
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# 单张图片使用示例(保留 ROI,不保存文件)
|
||||
# ---------------------------
|
||||
if __name__ == "__main__":
|
||||
model_path = r"best.pt"
|
||||
image_path = r"./test_image/1.jpg" # 单张图片路径
|
||||
roi_file = r"./roi_coordinates/1_rois.txt"
|
||||
target_size = 640
|
||||
threshold = 0.4 #加权得分阈值可以根据大小堆料分类结果进行调整
|
||||
|
||||
# 加载模型
|
||||
model = YOLO(model_path)
|
||||
|
||||
# 读取 ROI
|
||||
rois = load_global_rois(roi_file)
|
||||
if not rois:
|
||||
print("❌ 没有有效 ROI,退出")
|
||||
exit(1)
|
||||
|
||||
# 读取图片
|
||||
img = cv2.imread(image_path)
|
||||
if img is None:
|
||||
print(f"❌ 无法读取图片: {image_path}")
|
||||
exit(1)
|
||||
|
||||
# 注意:必须裁剪 ROI 并推理,因为训练的时候输入的图像是经过resize的
|
||||
crops = crop_and_resize(img, rois, target_size)
|
||||
for roi_resized, roi_idx in crops:
|
||||
#final_class, score, p1, p2 = classify_image_weighted(roi_resized, model, threshold=threshold)
|
||||
final_class,_,_,_ = classify_image_weighted(roi_resized, model, threshold=threshold)
|
||||
# 只输出信息,不保存文件
|
||||
#print(f"ROI {roi_idx} -> 类别: {final_class}, 加权分数: {score:.2f}, "
|
||||
#f"class1 置信度: {p1:.2f}, class2 置信度: {p2:.2f}")
|
||||
print(f"类别: {final_class}")
|
||||
|
||||
|
||||
1
zhuangtai_class_cls/roi_coordinates/1_rois1.txt
Normal file
1
zhuangtai_class_cls/roi_coordinates/1_rois1.txt
Normal file
@ -0,0 +1 @@
|
||||
626,725,1247,509
|
||||
BIN
zhuangtai_class_cls/test_image/1.jpg
Normal file
BIN
zhuangtai_class_cls/test_image/1.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 587 KiB |
BIN
zhuangtai_class_cls/test_image/2.jpg
Normal file
BIN
zhuangtai_class_cls/test_image/2.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 513 KiB |
BIN
zhuangtai_class_cls/test_image/3.jpg
Normal file
BIN
zhuangtai_class_cls/test_image/3.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 3.9 MiB |
198
zhuangtai_class_cls/tuili_f_yuantusave.py
Normal file
198
zhuangtai_class_cls/tuili_f_yuantusave.py
Normal file
@ -0,0 +1,198 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
import cv2
|
||||
import numpy as np
|
||||
from ultralytics import YOLO
|
||||
|
||||
# ---------------------------
|
||||
# 类别映射
|
||||
# ---------------------------
|
||||
CLASS_NAMES = {
|
||||
0: "未堆料",
|
||||
1: "小堆料",
|
||||
2: "大堆料",
|
||||
3: "未浇筑满",
|
||||
4: "浇筑满"
|
||||
}
|
||||
|
||||
# ---------------------------
|
||||
# 加载 ROI 列表
|
||||
# ---------------------------
|
||||
def load_global_rois(txt_path):
|
||||
rois = []
|
||||
if not os.path.exists(txt_path):
|
||||
print(f"❌ ROI 文件不存在: {txt_path}")
|
||||
return rois
|
||||
with open(txt_path, 'r') as f:
|
||||
for line in f:
|
||||
s = line.strip()
|
||||
if s:
|
||||
try:
|
||||
x, y, w, h = map(int, s.split(','))
|
||||
rois.append((x, y, w, h))
|
||||
except Exception as e:
|
||||
print(f"⚠️ 无法解析 ROI 行 '{s}': {e}")
|
||||
return rois
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# 裁剪并 resize ROI
|
||||
# ---------------------------
|
||||
def crop_and_resize(img, rois, target_size=640):
|
||||
crops = []
|
||||
h_img, w_img = img.shape[:2]
|
||||
for i, (x, y, w, h) in enumerate(rois):
|
||||
# 边界检查
|
||||
if x < 0 or y < 0 or x + w > w_img or y + h > h_img:
|
||||
print(f"⚠️ ROI 超出图像边界,跳过: ({x}, {y}, {w}, {h})")
|
||||
continue
|
||||
roi = img[y:y+h, x:x+w]
|
||||
roi_resized = cv2.resize(roi, (target_size, target_size), interpolation=cv2.INTER_AREA)
|
||||
crops.append((roi_resized, i))
|
||||
return crops
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# class1/class2 加权判断
|
||||
# ---------------------------
|
||||
def weighted_small_large(pred_probs, threshold=0.4, w1=0.3, w2=0.7):
|
||||
p1 = float(pred_probs[1])
|
||||
p2 = float(pred_probs[2])
|
||||
total = p1 + p2
|
||||
if total > 0:
|
||||
score = (w1 * p1 + w2 * p2) / total
|
||||
else:
|
||||
score = 0.0
|
||||
final_class = "大堆料" if score >= threshold else "小堆料"
|
||||
return final_class, score, p1, p2
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# 单张图片推理函数
|
||||
# ---------------------------
|
||||
def classify_image_weighted(image, model, threshold=0.5):
|
||||
results = model(image)
|
||||
pred_probs = results[0].probs.data.cpu().numpy().flatten()
|
||||
class_id = int(pred_probs.argmax())
|
||||
confidence = float(pred_probs[class_id])
|
||||
class_name = CLASS_NAMES.get(class_id, f"未知类别({class_id})")
|
||||
|
||||
# 对于 小堆料/大堆料 使用加权逻辑
|
||||
if class_id in [1, 2]:
|
||||
final_class, score, p1, p2 = weighted_small_large(pred_probs, threshold=threshold)
|
||||
else:
|
||||
final_class = class_name
|
||||
score = confidence
|
||||
p1 = float(pred_probs[1])
|
||||
p2 = float(pred_probs[2])
|
||||
|
||||
return final_class, score, p1, p2
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# 批量推理主函数(保存原图,不改名)
|
||||
# ---------------------------
|
||||
def batch_classify_images(model_path, input_folder, output_root, roi_file, target_size=640, threshold=0.5):
|
||||
# 加载模型
|
||||
print("🚀 加载模型...")
|
||||
model = YOLO(model_path)
|
||||
|
||||
# 确保输出根目录存在
|
||||
output_root = Path(output_root)
|
||||
output_root.mkdir(parents=True, exist_ok=True)
|
||||
print(f"📁 输出目录: {output_root}")
|
||||
|
||||
# 为所有类别创建子目录
|
||||
class_dirs = {}
|
||||
for name in CLASS_NAMES.values():
|
||||
d = output_root / name
|
||||
d.mkdir(exist_ok=True)
|
||||
class_dirs[name] = d
|
||||
|
||||
# 加载 ROI 区域
|
||||
rois = load_global_rois(roi_file)
|
||||
if not rois:
|
||||
print("❌ 没有有效 ROI,退出程序")
|
||||
return
|
||||
|
||||
print(f"🎯 已加载 {len(rois)} 个 ROI 区域")
|
||||
|
||||
# 定义缺陷严重性等级(数字越小越严重)
|
||||
severity_rank = {
|
||||
"未堆料": 0,
|
||||
"大堆料": 1,
|
||||
"小堆料": 2,
|
||||
"未浇筑满": 3,
|
||||
"浇筑满": 4
|
||||
}
|
||||
|
||||
# 遍历输入文件夹中的所有图片
|
||||
input_folder = Path(input_folder)
|
||||
image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff'}
|
||||
|
||||
processed_count = 0
|
||||
for img_path in input_folder.glob("*.*"):
|
||||
if img_path.suffix.lower() not in image_extensions:
|
||||
continue
|
||||
|
||||
try:
|
||||
print(f"\n📄 处理: {img_path.name}")
|
||||
img = cv2.imread(str(img_path))
|
||||
if img is None:
|
||||
print(f"❌ 无法读取图像: {img_path.name}")
|
||||
continue
|
||||
|
||||
# 裁剪并缩放 ROI
|
||||
crops = crop_and_resize(img, rois, target_size)
|
||||
|
||||
if not crops:
|
||||
print(f"⚠️ 无有效 ROI 裁剪区域: {img_path.name}")
|
||||
continue
|
||||
|
||||
detected_classes = []
|
||||
|
||||
# 遍历每个 ROI 进行分类
|
||||
for roi_img, roi_idx in crops:
|
||||
final_class, score, p1, p2 = classify_image_weighted(roi_img, model, threshold=threshold)
|
||||
detected_classes.append(final_class)
|
||||
print(f" 🔍 ROI{roi_idx}: {final_class} (score={score:.2f})")
|
||||
|
||||
# 选择最严重的类别
|
||||
most_severe_class = min(detected_classes, key=lambda x: severity_rank.get(x, 99))
|
||||
|
||||
# 构造目标路径:保持原文件名
|
||||
dst_path = class_dirs[most_severe_class] / img_path.name
|
||||
|
||||
# 保存原图(不修改内容,不重命名)
|
||||
cv2.imwrite(str(dst_path), img)
|
||||
print(f"✅ 保存 -> [{most_severe_class}] {dst_path}")
|
||||
|
||||
processed_count += 1
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ 处理失败 {img_path.name}: {e}")
|
||||
|
||||
print(f"\n🎉 批量处理完成!共处理 {processed_count} 张图像。")
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# 使用示例(请根据实际情况修改路径)
|
||||
# ---------------------------
|
||||
if __name__ == "__main__":
|
||||
# 🔧 用户配置区
|
||||
model_path = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/cls_resize/exp_cls2/weights/best.pt"
|
||||
input_folder = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1000"
|
||||
output_root = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1000/classified"
|
||||
roi_file = "./roi_coordinates/1_rois.txt"
|
||||
target_size = 640
|
||||
threshold = 0.4 # 小堆料 vs 大堆料的加权阈值
|
||||
|
||||
# 🚀 开始执行
|
||||
batch_classify_images(
|
||||
model_path=model_path,
|
||||
input_folder=input_folder,
|
||||
output_root=output_root,
|
||||
roi_file=roi_file,
|
||||
target_size=target_size,
|
||||
threshold=threshold
|
||||
)
|
||||
Reference in New Issue
Block a user