''' Author: [egrt] Date: 2022-08-23 11:44:15 LastEditors: Egrt LastEditTime: 2022-11-23 15:25:35 Description: HEAT的模型加载与预测 ''' from turtle import pos import torch import torch.nn as nn from models.resnet import ResNetBackbone from models.corner_models import HeatCorner from models.edge_models import HeatEdge from models.corner_to_edge import get_infer_edge_pairs from datasets.data_utils import get_pixel_features from huggingface_hub import hf_hub_download from PIL import Image from utils import image_utils from osgeo import gdal, ogr, osr from tqdm import tqdm import os import scipy import numpy as np import cv2 import skimage class HEAT(object): #-----------------------------------------# # 注意修改model_path #-----------------------------------------# _defaults = { #-----------------------------------------------# # model_data指向整体网络的地址 #-----------------------------------------------# "model_data" : 'model_data/heat_checkpoints/checkpoints/ckpts_heat_outdoor_256/checkpoint.pth', #-----------------------------------------------# # image_size模型预测图像的像素大小 #-----------------------------------------------# "image_size" : [256, 256], #-----------------------------------------------# # patch_size为模型切片的大小 #-----------------------------------------------# "patch_size" : 512, #-----------------------------------------------# # patch_overlap为切片重叠像素 #-----------------------------------------------# "patch_overlap" : 0, #-----------------------------------------------# # corner_thresh为预测角点的阈值大小 #-----------------------------------------------# "corner_thresh" : 0.01, #-----------------------------------------------# # 基于角点候选数的最大边数（不能大于6） #-----------------------------------------------# "corner_to_edge_multiplier": 3, #-----------------------------------------------# # 边缘推理筛选的迭代次数 #-----------------------------------------------# "infer_times" : 3, #-------------------------------# # 是否使用Cuda # 没有GPU可以设置成False #-------------------------------# "cuda" : False, } #---------------------------------------------------# # 初始化MASKGAN #---------------------------------------------------# def __init__(self, **kwargs): self.__dict__.update(self._defaults) for name, value in kwargs.items(): setattr(self, name, value) self.generate() def generate(self): # 从Huggingface加载整体网络模型 filepath = hf_hub_download(repo_id="Egrt/HEAT", filename="checkpoint.pth") self.model = torch.load(filepath) # 加载Backbone self.backbone = ResNetBackbone() strides = self.backbone.strides num_channels = self.backbone.num_channels self.backbone = nn.DataParallel(self.backbone) self.backbone = self.backbone.cuda() self.backbone.eval() # 加载角点检测模型 self.corner_model = HeatCorner(input_dim=128, hidden_dim=256, num_feature_levels=4, backbone_strides=strides, backbone_num_channels=num_channels) self.corner_model = nn.DataParallel(self.corner_model) self.corner_model = self.corner_model.cuda() self.corner_model.eval() # 加载边缘检测模型 self.edge_model = HeatEdge(input_dim=128, hidden_dim=256, num_feature_levels=4, backbone_strides=strides, backbone_num_channels=num_channels) self.edge_model = nn.DataParallel(self.edge_model) self.edge_model = self.edge_model.cuda() self.edge_model.eval() # 分别加载模型的地址 self.backbone.load_state_dict(self.model['backbone']) self.corner_model.load_state_dict(self.model['corner_model']) self.edge_model.load_state_dict(self.model['edge_model']) def detect_one_image(self, image): #---------------------------------------------------------# # 在这里将图像转换成RGB图像，防止灰度图在预测时报错。 # 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB #---------------------------------------------------------# image = cvtColor(image) # 这里判断图片是否需要分成多个patch if image.size[0] < self.patch_size or image.size[1] < self.patch_size: is_slice = False else: is_slice = True if is_slice: # 复制原图 image = np.array(image, dtype=np.uint8) # 复制输入的原图 viz_image = image.copy() height, width = image.shape[0], image.shape[1] # 获取缩放比例 scale = self.patch_size / self.image_size[0] # 初始化角点、边缘列表 pred_corners, pred_confs, pos_edges, edge_confs, c_outputs_np = [], [], [], [], [] # 开始切分 stride = self.patch_size - self.patch_overlap patch_boundingboxes = image_utils.compute_patch_boundingboxes((height, width), stride=stride, patch_res=self.patch_size) edge_len = 0 # 获取切分后的图片 for bbox in tqdm(patch_boundingboxes, desc="使用切分进行预测", leave=False): # 切分图像 crop_image = image[bbox[1]:bbox[3], bbox[0]:bbox[2], :] # np转Image类 crop_image = Image.fromarray(crop_image) try: pred_corners, pred_confs, pos_edges, edge_confs, c_outputs_np, _ = self.predict_no_patching(crop_image) except RuntimeError as e: print("ERROR: " + str(e)) print("INFO: 减小patch_size 直到适合内存") raise e # 拼接角点数组 pred_corners[:, 0] = pred_corners[:, 0] * scale + bbox[0] pred_corners[:, 1] = pred_corners[:, 1] * scale + bbox[1] pred_corners_viz = pred_corners viz_image = visualize_cond_generation(pred_corners_viz, pred_confs, viz_image, edges=pos_edges, edge_confs=edge_confs, shpfile=False) hr_image = Image.fromarray(np.uint8(viz_image)) else: pred_corners, pred_confs, pos_edges, edge_confs, c_outputs_np, viz_image = self.predict_no_patching(image) #---------------------------------------------------------# # 此处推理结束 # 开始在原图上根据角点坐标绘制角点与边缘 #---------------------------------------------------------# pred_corners_viz = pred_corners image_result = visualize_cond_generation(pred_corners_viz, pred_confs, viz_image, edges=pos_edges, edge_confs=edge_confs, shpfile=True) hr_image = Image.fromarray(np.uint8(image_result)) return hr_image #---------------------------------------------------------# # 不使用切片预测图像 # 返回预测后的角点坐标、边缘 #---------------------------------------------------------# def predict_no_patching(self, image): image = image.resize(tuple(self.image_size), Image.BICUBIC) # 将Image类转换为numpy image = np.array(image, dtype=np.uint8) # 复制输入的原图 viz_image = image.copy() # preprocess image numpy->tensor image = process_image(image) # 获取所有像素的位置编码, 默认的图像尺度为256 pixels, pixel_features = get_pixel_features(image_size=self.image_size[0]) # 开始模型的预测 with torch.no_grad(): image_feats, feat_mask, all_image_feats = self.backbone(image) pixel_features = pixel_features.unsqueeze(0).repeat(image.shape[0], 1, 1, 1) preds_s1 = self.corner_model(image_feats, feat_mask, pixel_features, pixels, all_image_feats) c_outputs = preds_s1 # 获取预测出的角点 c_outputs_np = c_outputs[0].detach().cpu().numpy() # 筛选出大于阈值的角点的坐标 pos_indices = np.where(c_outputs_np >= self.corner_thresh) pred_corners = pixels[pos_indices] # 获取对应预测角点的置信度 pred_confs = c_outputs_np[pos_indices] # 根据预测角点的置信度进行非极大抑制 pred_corners, pred_confs = corner_nms(pred_corners, pred_confs, image_size=c_outputs.shape[1]) # 对角点两两排列组合，获取所有的角点对 pred_corners, pred_confs, edge_coords, edge_mask, edge_ids = get_infer_edge_pairs(pred_corners, pred_confs) # 获取角点数量 corner_nums = torch.tensor([len(pred_corners)]).to(image.device) max_candidates = torch.stack([corner_nums.max() * self.corner_to_edge_multiplier] * len(corner_nums), dim=0) # 无序不重复集合 all_pos_ids = set() # 边缘置信度字典 all_edge_confs = dict() # 推理的迭代次数为3次 for tt in range(self.infer_times): if tt == 0: # gt_values和边缘掩膜大小一样且初始值为0 gt_values = torch.zeros_like(edge_mask).long() # 第一二维度的数值设置为2 gt_values[:, :] = 2 # 开始预测边缘 s1_logits, s2_logits_hb, s2_logits_rel, selected_ids, s2_mask, s2_gt_values = self.edge_model(image_feats, feat_mask,pixel_features,edge_coords, edge_mask,gt_values, corner_nums,max_candidates,True) num_total = s1_logits.shape[2] num_selected = selected_ids.shape[1] num_filtered = num_total - num_selected # 将输出值固定为(0,1)之间的概率分布 s1_preds = s1_logits.squeeze().softmax(0) s2_preds_rel = s2_logits_rel.squeeze().softmax(0) s2_preds_hb = s2_logits_hb.squeeze().softmax(0) s1_preds_np = s1_preds[1, :].detach().cpu().numpy() s2_preds_rel_np = s2_preds_rel[1, :].detach().cpu().numpy() s2_preds_hb_np = s2_preds_hb[1, :].detach().cpu().numpy() selected_ids = selected_ids.squeeze().detach().cpu().numpy() # 进行筛选，将(0.9, 1)之间的设置为T，将(0.01,0.9)之间的设置为U,(0,0.01)之间的设置为F if tt != self.infer_times - 1: s2_preds_np = s2_preds_hb_np pos_edge_ids = np.where(s2_preds_np >= 0.9) neg_edge_ids = np.where(s2_preds_np <= 0.01) for pos_id in pos_edge_ids[0]: actual_id = selected_ids[pos_id] if gt_values[0, actual_id] != 2: continue all_pos_ids.add(actual_id) all_edge_confs[actual_id] = s2_preds_np[pos_id] gt_values[0, actual_id] = 1 for neg_id in neg_edge_ids[0]: actual_id = selected_ids[neg_id] if gt_values[0, actual_id] != 2: continue gt_values[0, actual_id] = 0 num_to_pred = (gt_values == 2).sum() if num_to_pred <= num_filtered: break else: s2_preds_np = s2_preds_hb_np pos_edge_ids = np.where(s2_preds_np >= 0.5) for pos_id in pos_edge_ids[0]: actual_id = selected_ids[pos_id] if s2_mask[0][pos_id] is True or gt_values[0, actual_id] != 2: continue all_pos_ids.add(actual_id) all_edge_confs[actual_id] = s2_preds_np[pos_id] pos_edge_ids = list(all_pos_ids) edge_confs = [all_edge_confs[idx] for idx in pos_edge_ids] pos_edges = edge_ids[pos_edge_ids].cpu().numpy() edge_confs = np.array(edge_confs) if self.image_size[0] != 256: pred_corners = pred_corners / (self.image_size[0] / 256) return pred_corners, pred_confs, pos_edges, edge_confs, c_outputs_np, viz_image #---------------------------------------------------------# # 将图像转换成RGB图像，防止灰度图在预测时报错。 # 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB #---------------------------------------------------------# def cvtColor(image): if len(np.shape(image)) == 3 and np.shape(image)[2] == 3: return image else: image = image.convert('RGB') return image #---------------------------------------------------------# # 根据角点的置信度排序，并筛选出大于置信度的角点坐标 #---------------------------------------------------------# def corner_nms(preds, confs, image_size): data = np.zeros([image_size, image_size]) neighborhood_size = 5 threshold = 0 for i in range(len(preds)): data[preds[i, 1], preds[i, 0]] = confs[i] data_max = scipy.ndimage.filters.maximum_filter(data, neighborhood_size) maxima = (data == data_max) data_min = scipy.ndimage.filters.minimum_filter(data, neighborhood_size) diff = ((data_max - data_min) > threshold) maxima[diff == 0] = 0 results = np.where(maxima > 0) filtered_preds = np.stack([results[1], results[0]], axis=-1) new_confs = list() for i, pred in enumerate(filtered_preds): new_confs.append(data[pred[1], pred[0]]) new_confs = np.array(new_confs) return filtered_preds, new_confs def process_image(img): mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] img = skimage.img_as_float(img) img = img.transpose((2, 0, 1)) img = (img - np.array(mean)[:, np.newaxis, np.newaxis]) / np.array(std)[:, np.newaxis, np.newaxis] img = torch.Tensor(img).cuda() img = img.unsqueeze(0) return img def postprocess_preds(corners, confs, edges): corner_degrees = dict() for edge_i, edge_pair in enumerate(edges): corner_degrees[edge_pair[0]] = corner_degrees.setdefault(edge_pair[0], 0) + 1 corner_degrees[edge_pair[1]] = corner_degrees.setdefault(edge_pair[1], 0) + 1 good_ids = [i for i in range(len(corners)) if i in corner_degrees] if len(good_ids) == len(corners): return corners, confs, edges else: good_corners = corners[good_ids] good_confs = confs[good_ids] id_mapping = {value: idx for idx, value in enumerate(good_ids)} new_edges = list() for edge_pair in edges: new_pair = (id_mapping[edge_pair[0]], id_mapping[edge_pair[1]]) new_edges.append(new_pair) new_edges = np.array(new_edges) return good_corners, good_confs, new_edges #---------------------------------------------------------# # 将输入图像根据角点坐标进行可视化处理 # 不同于源代码，我们需要直接返回图像对象而不是保存到指定地址 #---------------------------------------------------------# def visualize_cond_generation(positive_pixels, confs, image, gt_corners=None, prec=None, recall=None, image_masks=None, edges=None, edge_confs=None, shpfile=False): # 复制原图 image = image.copy() if confs is not None: viz_confs = confs if edges is not None: preds = positive_pixels.astype(int) c_degrees = dict() for edge_i, edge_pair in enumerate(edges): conf = (edge_confs[edge_i] * 2) - 1 cv2.line(image, tuple(preds[edge_pair[0]]), tuple(preds[edge_pair[1]]), (255 * conf, 255 * conf, 0), 2) c_degrees[edge_pair[0]] = c_degrees.setdefault(edge_pair[0], 0) + 1 c_degrees[edge_pair[1]] = c_degrees.setdefault(edge_pair[1], 0) + 1 for idx, c in enumerate(positive_pixels): if edges is not None and idx not in c_degrees: continue if confs is None: cv2.circle(image, (int(c[0]), int(c[1])), 3, (0, 0, 255), -1) else: cv2.circle(image, (int(c[0]), int(c[1])), 3, (0, 0, 255 * viz_confs[idx]), -1) # if edges is not None: # cv2.putText(image, '{}'.format(c_degrees[idx]), (int(c[0]), int(c[1] - 5)), cv2.FONT_HERSHEY_SIMPLEX, # 0.5, (255, 0, 0), 1, cv2.LINE_AA) if gt_corners is not None: for c in gt_corners: cv2.circle(image, (int(c[0]), int(c[1])), 3, (0, 255, 0), -1) if image_masks is not None: mask_ids = np.where(image_masks == 1)[0] for mask_id in mask_ids: y_idx = mask_id // 64 x_idx = (mask_id - y_idx * 64) x_coord = x_idx * 4 y_coord = y_idx * 4 cv2.rectangle(image, (x_coord, y_coord), (x_coord + 3, y_coord + 3), (127, 127, 0), thickness=-1) # if confs is not None: # cv2.putText(image, 'max conf: {:.3f}'.format(confs.max()), (20, 20), cv2.FONT_HERSHEY_SIMPLEX, # 0.5, (255, 255, 0), 1, cv2.LINE_AA) if prec is not None: if isinstance(prec, tuple): cv2.putText(image, 'edge p={:.2f}, edge r={:.2f}'.format(prec[0], recall[0]), (20, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1, cv2.LINE_AA) cv2.putText(image, 'region p={:.2f}, region r={:.2f}'.format(prec[1], recall[1]), (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1, cv2.LINE_AA) else: cv2.putText(image, 'prec={:.2f}, recall={:.2f}'.format(prec, recall), (20, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1, cv2.LINE_AA) # 是否生成shp文件 if shpfile: preds = positive_pixels.astype(int) # 获取点列表 Polyline = [] for edge_i, edge_pair in enumerate(edges): Polyline.append([preds[edge_pair[0]], preds[edge_pair[1]]]) Polyline = np.array(Polyline, dtype=np.int32) # 写入shp文件 writeShp(save_file_dir="shpfile", Polyline=Polyline) return image def writeShp(save_file_dir="shpfile", Polyline=None): # 创建文件夹 if os.path.exists(save_file_dir) is False: os.makedirs(save_file_dir) # 支持中文路径 gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "YES") # 属性表字段支持中文 gdal.SetConfigOption("SHAPE_ENCODING", "UTF-8") # 注册驱动 ogr.RegisterAll() # 创建shp数据 strDriverName = "ESRI Shapefile" oDriver = ogr.GetDriverByName(strDriverName) if oDriver == None: return "驱动不可用："+strDriverName # 创建数据源 file_path = os.path.join(save_file_dir, "result.shp") oDS = oDriver.CreateDataSource(file_path) if oDS == None: return "创建文件失败：result.shp" if Polyline is not None: # 创建一个多边形图层，指定坐标系为WGS84 papszLCO = [] geosrs = osr.SpatialReference() geosrs.SetWellKnownGeogCS("WGS84") # 线：ogr_type = ogr.wkbLineString # 点：ogr_type = ogr.wkbPoint ogr_type = ogr.wkbMultiLineString # 面的类型为Polygon，线的类型为Polyline，点的类型为Point oLayer = oDS.CreateLayer("Polyline", geosrs, ogr_type, papszLCO) if oLayer == None: return "图层创建失败！" # 创建属性表 # 创建id字段 oId = ogr.FieldDefn("id", ogr.OFTInteger) oLayer.CreateField(oId, 1) # 创建name字段 oName = ogr.FieldDefn("name", ogr.OFTString) oLayer.CreateField(oName, 1) oDefn = oLayer.GetLayerDefn() # 创建要素 # 数据集 # wkt_geom id name point_str_list = ['({} {},{} {})'.format(row[0, 0], row[0, 1], row[1, 0], row[1, 1]) for row in Polyline] Polyline_Wkt = ','.join(point_str_list) features = ['Polyline0;MULTILINESTRING({})'.format(Polyline_Wkt)] for index, f in enumerate(features): oFeaturePolygon = ogr.Feature(oDefn) oFeaturePolygon.SetField("id",index) oFeaturePolygon.SetField("name",f.split(";")[0]) geomPolygon = ogr.CreateGeometryFromWkt(f.split(";")[1]) oFeaturePolygon.SetGeometry(geomPolygon) oLayer.CreateFeature(oFeaturePolygon) # 创建完成后，关闭进程 oDS.Destroy() return "数据集创建完成！"