模型文件路径支持配置,图片两边格子切掉
This commit is contained in:
liyubo
parent
085e1f4044
commit
d4151c2dab
@ -22,6 +22,7 @@ from middleware.minio_util import upload_file, downFile, check_zip_size, upload_
|
||||
from util.yolo2pix_new import *
|
||||
from util.smb_tool import *
|
||||
import threading
|
||||
import json
|
||||
|
||||
# 定义红白蓝颜色 (BGR格式)
|
||||
RED = (0, 0, 255)
|
||||
@ -653,9 +654,15 @@ class YOLOSegmentationInference:
|
||||
|
||||
# 获取当前脚本文件所在的目录
|
||||
script_dir = os.path.dirname(__file__)
|
||||
# 构建模型文件的绝对路径
|
||||
asphalt_model_path = os.path.abspath(os.path.join(script_dir, '..', 'pt_save', 'asphalt20260212_0.391.pt'))
|
||||
cream_model_path = os.path.abspath(os.path.join(script_dir, '..', 'pt_save', 'asphalt20260212_0.391.pt'))
|
||||
# 获取模型json配置文件
|
||||
with open(os.path.join(script_dir, 'model.json'), 'r') as f:
|
||||
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 :
|
||||
road_type_cn = road_data['路面类型（沥青/水泥/砂石）']
|
||||
@ -1020,4 +1027,4 @@ def predict_images_share_dir(task_id, pt_name, zip_url, user_name, pwd, output_d
|
||||
else :
|
||||
print(f"错误: 输入 {zip_url} 不是有效的文件或目录")
|
||||
|
||||
return standardized_path(f"{target_path}_识别/{task_id}/{current_time}/{task_id}.zip"), "success"
|
||||
return standardized_path(f"{target_path}_识别/{task_id}/{current_time}/{task_id}.zip"), "success"
|
||||
@ -809,11 +809,11 @@ def get_road_dict(dir,user_name,pwd) :
|
||||
|
||||
found_paths = scanner.find_files_by_name(
|
||||
share_path=config['share'],
|
||||
file_name='每公里指标明细表*.xls',
|
||||
file_name='每公里指标明细表*.xls*',
|
||||
start_dir=config['dir'],
|
||||
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):
|
||||
print(f"{i}. {path}")
|
||||
|
||||
|
||||
@ -22,6 +22,7 @@ CELL_AREA = 0.01 # 每格面积 (平方米)
|
||||
GRID_WIDTH = 108 # 网格像素宽
|
||||
GRID_HEIGHT = 102 # 网格像素高
|
||||
COVER_RATIO = 0.01 # mask 覆盖比例阈值
|
||||
RESERVED_CELL_NUM = 20 # 道路灾害格子过滤掉两侧数据，只保留每行中间20个格子
|
||||
|
||||
# ---------------- 路面类别映射 ----------------
|
||||
CLASS_MAP_ASPHALT = {
|
||||
@ -160,7 +161,8 @@ def detect_road_type_from_road_dict(road_dict, pile_dict, img_file_name):
|
||||
road_type = "gravel"
|
||||
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标签转换成格子编号和类别"""
|
||||
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)
|
||||
@ -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)
|
||||
if img is None: return "", {}
|
||||
h, w = img.shape[:2]
|
||||
cols = max(1, w//GRID_WIDTH)
|
||||
rows = max(1, h//GRID_HEIGHT)
|
||||
cols = max(1, w//grid_width)
|
||||
rows = max(1, h//grid_height)
|
||||
|
||||
result_lines = []
|
||||
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
|
||||
for r in range(rows):
|
||||
for c in range(cols):
|
||||
x1,y1 = c*GRID_WIDTH, r*GRID_HEIGHT
|
||||
x2,y2 = min(w,x1+GRID_WIDTH), min(h,y1+GRID_HEIGHT)
|
||||
x1,y1 = c*grid_width, r*grid_height
|
||||
x2,y2 = min(w,x1+grid_width), min(h,y1+grid_height)
|
||||
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)
|
||||
# 最小x坐标，y坐标
|
||||
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
|
||||
|
||||
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):
|
||||
"""将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):
|
||||
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):
|
||||
print(f"⚠️ 找不到标签: {label_file}")
|
||||
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_path = os.path.splitext(image_path)[0]+"_grid.txt"
|
||||
with open(grid_txt_path,'w',encoding='utf-8') as f:
|
||||
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_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: 路线编码到路线信息的映射字典
|
||||
"""
|
||||
# 查找匹配的Excel文件
|
||||
pattern_xls = os.path.join(local_dir, '每公里指标明细表*.xls')
|
||||
pattern_xlsx = os.path.join(local_dir, '每公里指标明细表*.xlsx')
|
||||
found_paths = glob.glob(pattern_xls) + glob.glob(pattern_xlsx)
|
||||
pattern = os.path.join(local_dir, '每公里指标明细表*.xls*')
|
||||
found_paths = glob.glob(pattern)
|
||||
|
||||
print(f"\n找到 {len(found_paths)} 个 '每公里指标明细表*.xls/xlsx' 文件:")
|
||||
print(f"\n找到 {len(found_paths)} 个 '每公里指标明细表*.xls*' 文件:")
|
||||
for i, path in enumerate(found_paths, 1):
|
||||
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)
|
||||
# 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)
|
||||
road_dict = get_road_dict(output_dir)
|
||||
process_dir(road_dict, pile_dict, output_dir)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user