174 lines
5.8 KiB
Python
174 lines
5.8 KiB
Python
|
import os.path
|
|||
|
|
|
|||
|
|
import numpy as np
|
|||
|
|
import cv2
|
|||
|
|
|
|||
|
|
# 根据输出结果输出json问价还有边界
|
|||
|
|
def segment_instance_and_contour(seg_map, class_names):
|
|||
|
|
"""
|
|||
|
|
参数：
|
|||
|
|
seg_map: HxW 的语义分割图像，像素值表示类别编号
|
|||
|
|
class_names: 类别编号到类别名的映射，如 {0:'Background', 1:'Farmland', 2:'Forest'}
|
|||
|
|
返回：
|
|||
|
|
实例列表，每个包含id、类别名、轮廓点
|
|||
|
|
"""
|
|||
|
|
instance_id = 1
|
|||
|
|
results = []
|
|||
|
|
|
|||
|
|
for cls_id in np.unique(seg_map):
|
|||
|
|
if cls_id == 0: # 忽略背景
|
|||
|
|
continue
|
|||
|
|
mask = (seg_map == cls_id).astype(np.uint8)
|
|||
|
|
|
|||
|
|
# 找连通域（不同地块）
|
|||
|
|
num_labels, labels = cv2.connectedComponents(mask)
|
|||
|
|
|
|||
|
|
for i in range(1, num_labels): # 第0是背景
|
|||
|
|
instance_mask = (labels == i).astype(np.uint8)
|
|||
|
|
|
|||
|
|
# 找轮廓
|
|||
|
|
contours, _ = cv2.findContours(instance_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|||
|
|
for contour in contours:
|
|||
|
|
contour_list = contour.squeeze().tolist()
|
|||
|
|
if len(contour_list) < 3:
|
|||
|
|
continue # 过滤掉太小的碎片
|
|||
|
|
results.append({
|
|||
|
|
'id': instance_id,
|
|||
|
|
'class': class_names.get(cls_id, f"Class{cls_id}"),
|
|||
|
|
'contour': contour_list
|
|||
|
|
})
|
|||
|
|
instance_id += 1
|
|||
|
|
|
|||
|
|
return results
|
|||
|
|
|
|||
|
|
def draw_instances_with_ids(seg_map, results, save_path="instance_visual.png"):
|
|||
|
|
"""
|
|||
|
|
在原始图或灰度图上绘制每个地块的轮廓及其编号。
|
|||
|
|
|
|||
|
|
参数：
|
|||
|
|
seg_map: 原始图像或分割图（HxW 或 HxWx3）
|
|||
|
|
results: 上一步输出的地块信息列表
|
|||
|
|
save_path: 保存路径
|
|||
|
|
"""
|
|||
|
|
if len(seg_map.shape) == 2: # 若为灰度图，转为BGR
|
|||
|
|
vis_img = cv2.cvtColor(seg_map, cv2.COLOR_GRAY2BGR)
|
|||
|
|
else:
|
|||
|
|
vis_img = seg_map.copy()
|
|||
|
|
# h, w = vis_img.shape[:2]
|
|||
|
|
for item in results:
|
|||
|
|
contour = np.array(item['contour']).reshape(-1, 1, 2).astype(np.int32)
|
|||
|
|
|
|||
|
|
# 随机颜色（根据类别也可指定颜色）
|
|||
|
|
color = tuple(np.random.randint(0, 255, size=3).tolist())
|
|||
|
|
cv2.drawContours(vis_img, [contour], -1, color, 2)#线的宽度
|
|||
|
|
|
|||
|
|
# 计算重心位置以便放编号
|
|||
|
|
M = cv2.moments(contour)
|
|||
|
|
if M['m00'] != 0:
|
|||
|
|
cx = int(M['m10'] / M['m00'])
|
|||
|
|
cy = int(M['m01'] / M['m00'])
|
|||
|
|
cv2.putText(vis_img, str(item['id']), (cx, cy),
|
|||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, cv2.LINE_AA)#序号的大小与线宽度
|
|||
|
|
|
|||
|
|
# # 替代 moments 重心方式，使用外接矩形中心
|
|||
|
|
# x, y, w, h = cv2.boundingRect(contour)
|
|||
|
|
# cx = x + w // 2
|
|||
|
|
# cy = y + h // 2
|
|||
|
|
# cv2.putText(vis_img, str(item['id']), (cx, cy),
|
|||
|
|
# cv2.FONT_HERSHEY_SIMPLEX, 3, color, 2, cv2.LINE_AA)
|
|||
|
|
|
|||
|
|
# # 创建该块的 instance mask，用于质心计算
|
|||
|
|
# instance_mask = np.zeros((h, w), dtype=np.uint8)
|
|||
|
|
# cv2.drawContours(instance_mask, [contour], -1, 1, -1)
|
|||
|
|
#
|
|||
|
|
# # 计算质心（内部像素点的几何中心）
|
|||
|
|
# ys, xs = np.where(instance_mask > 0)
|
|||
|
|
# if len(xs) > 0 and len(ys) > 0:
|
|||
|
|
# cx = int(np.mean(xs))
|
|||
|
|
# cy = int(np.mean(ys))
|
|||
|
|
# cv2.putText(vis_img, str(item['id']), (cx, cy),
|
|||
|
|
# cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, cv2.LINE_AA)
|
|||
|
|
|
|||
|
|
cv2.imwrite(save_path, vis_img)
|
|||
|
|
print(f"✅ 可视化图保存至: {save_path}")
|
|||
|
|
|
|||
|
|
def color_map_to_label(seg_img, palette):
|
|||
|
|
"""
|
|||
|
|
将彩色语义图映射为类别ID图
|
|||
|
|
|
|||
|
|
参数：
|
|||
|
|
seg_img: HxWx3 彩色图像（例如读取的png）
|
|||
|
|
palette: 类别调色板，如 [[0,0,0], [139,69,19], ...]
|
|||
|
|
|
|||
|
|
返回：
|
|||
|
|
label_map: HxW 图像，每个像素为类别ID（int）
|
|||
|
|
"""
|
|||
|
|
h, w, _ = seg_img.shape
|
|||
|
|
label_map = np.zeros((h, w), dtype=np.uint8)
|
|||
|
|
|
|||
|
|
# 将调色板转换为字典
|
|||
|
|
color2label = {tuple(color): idx for idx, color in enumerate(palette)}
|
|||
|
|
|
|||
|
|
# 对每个像素映射
|
|||
|
|
for color, label in color2label.items():
|
|||
|
|
mask = np.all(seg_img == color, axis=-1)
|
|||
|
|
label_map[mask] = label
|
|||
|
|
|
|||
|
|
return label_map
|
|||
|
|
|
|||
|
|
# 定义类别调色板和类别名
|
|||
|
|
PALETTE = [
|
|||
|
|
[0, 0, 0], # background
|
|||
|
|
[139, 69, 19], # barren
|
|||
|
|
[0, 255, 0], # forest
|
|||
|
|
[255, 255, 0], # farmland
|
|||
|
|
[0, 0, 255], # water
|
|||
|
|
[128, 128, 128], # road
|
|||
|
|
[0, 255, 255] # building
|
|||
|
|
]
|
|||
|
|
CLASS_NAMES = {
|
|||
|
|
0: 'background',
|
|||
|
|
1: 'barren',
|
|||
|
|
2: 'forest',
|
|||
|
|
3: 'farmland',
|
|||
|
|
4: 'water',
|
|||
|
|
5: 'road',
|
|||
|
|
6: 'building'
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def draw_json_boundary(pic_url):
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 读取彩色分割图
|
|||
|
|
seg_color = cv2.imread(pic_url) # 注意是 BGR 格式
|
|||
|
|
# seg_color = cv2.imread(r"D:\project\UAV_model\prediction_results\unetformer_UAV_6000X4000\patch_0028.png") # 注意是 BGR 格式
|
|||
|
|
# seg_color = cv2.imread(r"J:\uhr\pytools\numer_coordinate\patch_0047_rgb.png") # 注意是 BGR 格式
|
|||
|
|
seg_color = cv2.cvtColor(seg_color, cv2.COLOR_BGR2RGB)
|
|||
|
|
|
|||
|
|
# 转为类别索引图
|
|||
|
|
label_map = color_map_to_label(seg_color, PALETTE)
|
|||
|
|
|
|||
|
|
# 继续后续处理
|
|||
|
|
results = segment_instance_and_contour(label_map, CLASS_NAMES)
|
|||
|
|
dir_name=os.path.dirname(pic_url)
|
|||
|
|
final_vis_png_name=os.path.basename(pic_url)+"final_vis.png"
|
|||
|
|
final_vis_png_path=os.path.join(dir_name,final_vis_png_name) #输出大本地的只有边缘线条
|
|||
|
|
|
|||
|
|
instance_results_json_name=os.path.basename(pic_url)+"instance_results.json"
|
|||
|
|
instance_results_json_path=os.path.join(dir_name,instance_results_json_name) # 输出到本地的，分割的边缘的json文件
|
|||
|
|
draw_instances_with_ids(label_map, results, save_path=final_vis_png_path)
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 保存为 JSON（可选）
|
|||
|
|
import json
|
|||
|
|
with open(instance_results_json_path, "w") as f:
|
|||
|
|
json.dump(results, f, indent=2)
|
|||
|
|
|
|||
|
|
return final_vis_png_path,instance_results_json_path
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|