bdzl2/ai-train_platform
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

模型文件路径支持配置,图片两边格子切掉

Browse Source
This commit is contained in:
liyubo 2026-03-05 09:21:07 +08:00
parent 085e1f4044
commit d4151c2dab
3 changed files with 34 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
predict/predict_yolo11seg.py
View File

@ -22,6 +22,7 @@ from middleware.minio_util import upload_file, downFile, check_zip_size, upload_
from util.yolo2pix_new import * from util.yolo2pix_new import *
from util.smb_tool import * from util.smb_tool import *
import threading import threading
import json
# 定义红白蓝颜色 (BGR格式) # 定义红白蓝颜色 (BGR格式)
RED = (0, 0, 255) RED = (0, 0, 255)
@ -653,9 +654,15 @@ class YOLOSegmentationInference:
# 获取当前脚本文件所在的目录 # 获取当前脚本文件所在的目录
script_dir = os.path.dirname(__file__) script_dir = os.path.dirname(__file__)
# 构建模型文件的绝对路径 # 获取模型json配置文件
asphalt_model_path = os.path.abspath(os.path.join(script_dir, '..', 'pt_save', 'asphalt20260212_0.391.pt')) with open(os.path.join(script_dir, 'model.json'), 'r') as f:
cream_model_path = os.path.abspath(os.path.join(script_dir, '..', 'pt_save', 'asphalt20260212_0.391.pt')) config = json.load(f)
# 构建模型文件的绝对路径
asphalt_model_path = os.path.abspath(os.path.join(script_dir, '..', 'pt_save', config['asphalt']))
cream_model_path = os.path.abspath(os.path.join(script_dir, '..', 'pt_save', config['cream']))
print(f"asphalt_model_path={asphalt_model_path}")
print(f"cream_model_path={cream_model_path}")
if road_data : if road_data :
road_type_cn = road_data['路面类型（沥青/水泥/砂石）'] road_type_cn = road_data['路面类型（沥青/水泥/砂石）']

4
util/smb_tool.py
View File

@ -809,11 +809,11 @@ def get_road_dict(dir,user_name,pwd) :
found_paths = scanner.find_files_by_name( found_paths = scanner.find_files_by_name(
share_path=config['share'], share_path=config['share'],
file_name='每公里指标明细表*.xls', file_name='每公里指标明细表*.xls*',
start_dir=config['dir'], start_dir=config['dir'],
max_depth=4 max_depth=4
) )
print(f"\n找到 {len(found_paths)} 个 '每公里指标明细表*.xls' 文件:") print(f"\n找到 {len(found_paths)} 个 '每公里指标明细表*.xls*' 文件:")
for i, path in enumerate(found_paths, 1): for i, path in enumerate(found_paths, 1):
print(f"{i}. {path}") print(f"{i}. {path}")

35
util/yolo2pix_new.py
View File

@ -22,6 +22,7 @@ CELL_AREA = 0.01 # 每格面积 (平方米)
GRID_WIDTH = 108 # 网格像素宽 GRID_WIDTH = 108 # 网格像素宽
GRID_HEIGHT = 102 # 网格像素高 GRID_HEIGHT = 102 # 网格像素高
COVER_RATIO = 0.01 # mask 覆盖比例阈值 COVER_RATIO = 0.01 # mask 覆盖比例阈值
RESERVED_CELL_NUM = 20 # 道路灾害格子过滤掉两侧数据，只保留每行中间20个格子
# ---------------- 路面类别映射 ---------------- # ---------------- 路面类别映射 ----------------
CLASS_MAP_ASPHALT = { CLASS_MAP_ASPHALT = {
@ -160,7 +161,8 @@ def detect_road_type_from_road_dict(road_dict, pile_dict, img_file_name):
road_type = "gravel" road_type = "gravel"
return road_code, pile_no, road_type return road_code, pile_no, road_type
def yoloseg_to_grid_share_dir(road_dict,pile_dict,image_path,label_file,grid_width=GRID_WIDTH,grid_height=GRID_HEIGHT,cover_ratio=COVER_RATIO): def yoloseg_to_grid_share_dir(road_dict,pile_dict,image_path,label_file,cell_index_start,cell_index_end,
grid_width=GRID_WIDTH,grid_height=GRID_HEIGHT,cover_ratio=COVER_RATIO):
"""将YOLO-Seg标签转换成格子编号和类别""" """将YOLO-Seg标签转换成格子编号和类别"""
img_file_name = os.path.basename(image_path) img_file_name = os.path.basename(image_path)
road_code, pile_no, road_type = detect_road_type_from_road_dict(road_dict, pile_dict, img_file_name) road_code, pile_no, road_type = detect_road_type_from_road_dict(road_dict, pile_dict, img_file_name)
@ -172,8 +174,8 @@ def yoloseg_to_grid_share_dir(road_dict,pile_dict,image_path,label_file,grid_wid
img = cv2.imread(image_path) img = cv2.imread(image_path)
if img is None: return "", {} if img is None: return "", {}
h, w = img.shape[:2] h, w = img.shape[:2]
cols = max(1, w//GRID_WIDTH) cols = max(1, w//grid_width)
rows = max(1, h//GRID_HEIGHT) rows = max(1, h//grid_height)
result_lines = [] result_lines = []
all_class_cells = {} all_class_cells = {}
@ -198,10 +200,10 @@ def yoloseg_to_grid_share_dir(road_dict,pile_dict,image_path,label_file,grid_wid
max_y = 0 max_y = 0
for r in range(rows): for r in range(rows):
for c in range(cols): for c in range(cols):
x1,y1 = c*GRID_WIDTH, r*GRID_HEIGHT x1,y1 = c*grid_width, r*grid_height
x2,y2 = min(w,x1+GRID_WIDTH), min(h,y1+GRID_HEIGHT) x2,y2 = min(w,x1+grid_width), min(h,y1+grid_height)
region = mask[y1:y2, x1:x2] region = mask[y1:y2, x1:x2]
if np.count_nonzero(region)/region.size>cover_ratio: if c >= cell_index_start and c < cell_index_end and np.count_nonzero(region)/region.size>cover_ratio:
covered_cells.append(r*cols+c+1) covered_cells.append(r*cols+c+1)
# 最小x坐标，y坐标 # 最小x坐标，y坐标
if min_x > c : if min_x > c :
@ -236,7 +238,7 @@ def yoloseg_to_grid_share_dir(road_dict,pile_dict,image_path,label_file,grid_wid
}) })
all_class_cells[cname] = cell_info all_class_cells[cname] = cell_info
return '\n'.join(result_lines), all_class_cells, road_type, rows * cols return '\n'.join(result_lines), all_class_cells, road_type, rows * RESERVED_CELL_NUM
def yoloseg_to_grid(image_path,label_file,cover_ratio=COVER_RATIO): def yoloseg_to_grid(image_path,label_file,cover_ratio=COVER_RATIO):
"""将YOLO-Seg标签转换成格子编号和类别""" """将YOLO-Seg标签转换成格子编号和类别"""
@ -556,6 +558,12 @@ def format_number_to_k_code(number):
# ---------------- 主函数-共享目录 ---------------- # ---------------- 主函数-共享目录 ----------------
def process_dir(road_dict,pile_dict,dir="output",cell_area=CELL_AREA,cell_width=CELL_WIDTH,cell_height=CELL_HEIGHT,grid_width=GRID_WIDTH,grid_height=GRID_HEIGHT): def process_dir(road_dict,pile_dict,dir="output",cell_area=CELL_AREA,cell_width=CELL_WIDTH,cell_height=CELL_HEIGHT,grid_width=GRID_WIDTH,grid_height=GRID_HEIGHT):
os.makedirs(dir,exist_ok=True) os.makedirs(dir,exist_ok=True)
# 网格像素宽
road_recognize_width = road_dict.get('识别宽度（米）', 3.6)
grid_width = round(4096 * CELL_WIDTH / road_recognize_width)
# 道路灾害格子过滤掉两侧数据，只保留每行中间20个格子
cell_index_start = round((road_recognize_width / CELL_WIDTH - RESERVED_CELL_NUM) / 2)
cell_index_end = cell_index_start + RESERVED_CELL_NUM
# 解压 # 解压
# 读取桩号映射 # 读取桩号映射
# 遍历图片 # 遍历图片
@ -569,13 +577,13 @@ def process_dir(road_dict,pile_dict,dir="output",cell_area=CELL_AREA,cell_width=
if not os.path.exists(label_file): if not os.path.exists(label_file):
print(f"⚠️ 找不到标签: {label_file}") print(f"⚠️ 找不到标签: {label_file}")
continue continue
out_txt, class_cells, road_type, all_cell_num = yoloseg_to_grid_share_dir(road_dict,pile_dict,image_path,label_file) out_txt, class_cells, road_type, all_cell_num = yoloseg_to_grid_share_dir(road_dict,pile_dict,image_path,label_file,cell_index_start,cell_index_end,grid_width=grid_width)
# 写每张图独立 _grid.txt # 写每张图独立 _grid.txt
grid_txt_path = os.path.splitext(image_path)[0]+"_grid.txt" grid_txt_path = os.path.splitext(image_path)[0]+"_grid.txt"
with open(grid_txt_path,'w',encoding='utf-8') as f: with open(grid_txt_path,'w',encoding='utf-8') as f:
f.write(out_txt) f.write(out_txt)
# 生成网格可视化 # 生成网格可视化
draw_grid_on_image(image_path,class_cells,cell_size=(GRID_WIDTH, GRID_HEIGHT),save_path=os.path.splitext(image_path)[0]+"_grid.jpg") draw_grid_on_image(image_path,class_cells,cell_size=(grid_width, grid_height),save_path=os.path.splitext(image_path)[0]+"_grid.jpg")
# 统计各类面积 # 统计各类面积
# counts = {k:[len(v[0])*cell_area, v[1][0], v[1][1]] for k,v in class_cells.items()} # counts = {k:[len(v[0])*cell_area, v[1][0], v[1][1]] for k,v in class_cells.items()}
counts = {k:[[len(v_child.get("covered_cells"))*cell_area, v_child.get("width"), v_child.get("height")] for v_child in v] for k,v in class_cells.items()} counts = {k:[[len(v_child.get("covered_cells"))*cell_area, v_child.get("width"), v_child.get("height")] for v_child in v] for k,v in class_cells.items()}
@ -906,11 +914,10 @@ def get_road_dict(local_dir):
dict: 路线编码到路线信息的映射字典 dict: 路线编码到路线信息的映射字典
""" """
# 查找匹配的Excel文件 # 查找匹配的Excel文件
pattern_xls = os.path.join(local_dir, '每公里指标明细表*.xls') pattern = os.path.join(local_dir, '每公里指标明细表*.xls*')
pattern_xlsx = os.path.join(local_dir, '每公里指标明细表*.xlsx') found_paths = glob.glob(pattern)
found_paths = glob.glob(pattern_xls) + glob.glob(pattern_xlsx)
print(f"\n找到 {len(found_paths)} 个 '每公里指标明细表*.xls/xlsx' 文件:") print(f"\n找到 {len(found_paths)} 个 '每公里指标明细表*.xls*' 文件:")
for i, path in enumerate(found_paths, 1): for i, path in enumerate(found_paths, 1):
print(f"{i}. {path}") print(f"{i}. {path}")
@ -989,7 +996,7 @@ if __name__=="__main__":
# calc_cell_area, calc_grid_width, calc_grid_height = calc_grid_param(2048, 4096, 3.6, 2) # calc_cell_area, calc_grid_width, calc_grid_height = calc_grid_param(2048, 4096, 3.6, 2)
# print(f"calc_cell_area={calc_cell_area}, calc_grid_width={calc_grid_width}, calc_grid_height={calc_grid_height}") # print(f"calc_cell_area={calc_cell_area}, calc_grid_width={calc_grid_width}, calc_grid_height={calc_grid_height}")
output_dir = "D:/devForBdzlWork/ai-train_platform/predictions/jlp/20260227090742" output_dir = "D:/devForBdzlWork/ai-train_platform/predictions/jlp/C006500107A"
pile_dict = get_pile_dict(output_dir) pile_dict = get_pile_dict(output_dir)
road_dict = get_road_dict(output_dir) road_dict = get_road_dict(output_dir)
process_dir(road_dict, pile_dict, output_dir) process_dir(road_dict, pile_dict, output_dir)