Hugging Face's logo

Spaces:
Egrt
/
HEAT
Runtime error

App Files Community
HEAT / HEAT.py
Egrt's picture
Egrt
init
424188c
raw history blame contribute delete
No virus
21.6 kB
'''
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 "数据集创建完成！"