389 lines
17 KiB
Python
389 lines
17 KiB
Python
|
import numpy as np
|
|||
|
|
import copy
|
|||
|
|
import cv2 as cv
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
import os
|
|||
|
|
import torch
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
from PIL import Image
|
|||
|
|
import uuid
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
class segmentation():
|
|||
|
|
|
|||
|
|
def __init__(self, model_name="xx"):
|
|||
|
|
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|||
|
|
self.base_dir = os.path.dirname(os.path.abspath(__file__))
|
|||
|
|
self.model_name = model_name
|
|||
|
|
self.palette = [0, 0, 0, 255, 0, 0, 0, 255, 0]
|
|||
|
|
self.fusin = True
|
|||
|
|
|
|||
|
|
def create_tmp_path(self, path):
|
|||
|
|
folder_name = str(uuid.uuid4())
|
|||
|
|
folder_path = os.path.join(path,'tmp', folder_name)
|
|||
|
|
os.makedirs(folder_path)
|
|||
|
|
return folder_path
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 图片分割识别结果
|
|||
|
|
def start_segmentation(self,path,net):
|
|||
|
|
assert os.path.isfile(path)
|
|||
|
|
image = Image.open(path)
|
|||
|
|
orininal_h = image.size[1]
|
|||
|
|
orininal_w = image.size[0]
|
|||
|
|
item = {"orininal_h": orininal_h, "orininal_w": orininal_w, "image": image}
|
|||
|
|
image = item["image"]
|
|||
|
|
orininal_w = item["orininal_w"]
|
|||
|
|
orininal_h = item["orininal_h"]
|
|||
|
|
old_img = copy.deepcopy(image)
|
|||
|
|
imaged = cv.resize(np.array(image), dsize=(512, 512), interpolation=cv.INTER_LINEAR)
|
|||
|
|
image_data = np.expand_dims(
|
|||
|
|
np.transpose(self.preprocess_input(np.array(imaged, np.float32), md=False), (2, 0, 1)), 0)
|
|||
|
|
|
|||
|
|
with torch.no_grad():
|
|||
|
|
images = torch.from_numpy(image_data)
|
|||
|
|
image = images.to(device="cuda", dtype=torch.float32)
|
|||
|
|
model = net.to(device="cuda")
|
|||
|
|
out = model(image) # batch_size, 2, 512, 512
|
|||
|
|
if isinstance(out, list) or isinstance(out, tuple): # 可能有多个输出（这里把辅助解码头的也输出的所以是多个）
|
|||
|
|
out = out[0] # 就取第一个
|
|||
|
|
out = out[0] # 去掉batch
|
|||
|
|
pr = F.softmax(out.permute(1, 2, 0), dim=-1).cpu().numpy()
|
|||
|
|
result = pr.argmax(axis=-1)
|
|||
|
|
# 二值图像
|
|||
|
|
output_im = Image.fromarray(np.uint8(result)).convert('P')
|
|||
|
|
output_im.putpalette(self.palette)
|
|||
|
|
file_name = os.path.basename(path)
|
|||
|
|
file_directory = os.path.dirname(path)
|
|||
|
|
file_name_ = "binary_"+ file_name
|
|||
|
|
file_path = os.path.join(file_directory,file_name_)
|
|||
|
|
output_im.save(file_path)
|
|||
|
|
# 融合图像
|
|||
|
|
PALETTE = [(0, 0, 0), (255, 0, 0), (0, 255, 0)] # 这里主动改分类了，将识别图像与底图融合
|
|||
|
|
|
|||
|
|
seg_img0 = np.reshape(np.array(PALETTE, np.uint8)[np.reshape(result, [-1])],
|
|||
|
|
[orininal_h, orininal_w, -1])
|
|||
|
|
|
|||
|
|
image0 = Image.fromarray(np.uint8(seg_img0))
|
|||
|
|
|
|||
|
|
fusion = Image.blend(old_img, image0, 0.4) # 图像融合
|
|||
|
|
file_name_ = "segment_"+ file_name
|
|||
|
|
file_path = os.path.join(file_directory,file_name_)
|
|||
|
|
fusion.save(file_path)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 获取大型图片，需要进行裁剪
|
|||
|
|
def cut_big_image(self,path):
|
|||
|
|
file_name = os.path.basename(path)
|
|||
|
|
file_directory = os.path.dirname(path)
|
|||
|
|
output_floder_normal = os.path.join(file_directory,'ori')
|
|||
|
|
# 创建输出文件夹（如果不存在）
|
|||
|
|
os.makedirs(output_floder_normal, exist_ok=True)
|
|||
|
|
# 打开图像
|
|||
|
|
image = Image.open(path)
|
|||
|
|
width, height = image.size
|
|||
|
|
# 定义块的大小
|
|||
|
|
block_size = 512
|
|||
|
|
|
|||
|
|
# 计算需要的块数
|
|||
|
|
num_blocks_x = (width + block_size - 1) // block_size # 向上取整
|
|||
|
|
num_blocks_y = (height + block_size - 1) // block_size # 向上取整
|
|||
|
|
# 裁剪并保存图像块
|
|||
|
|
for i in range(num_blocks_x):
|
|||
|
|
for j in range(num_blocks_y):
|
|||
|
|
# 计算裁剪区域
|
|||
|
|
left = i * block_size
|
|||
|
|
upper = j * block_size
|
|||
|
|
right = min(left + block_size, width)
|
|||
|
|
lower = min(upper + block_size, height)
|
|||
|
|
|
|||
|
|
# 裁剪图像
|
|||
|
|
block = image.crop((left, upper, right, lower))
|
|||
|
|
|
|||
|
|
# 创建一个新的512x512白色图像
|
|||
|
|
new_block = Image.new('RGB', (block_size, block_size), (255, 255, 255))
|
|||
|
|
# 将裁剪的图像粘贴到白色图像上
|
|||
|
|
new_block.paste(block, (0, 0))
|
|||
|
|
|
|||
|
|
# 保存图像块
|
|||
|
|
block_filename = f'block_{i}_{j}.png'
|
|||
|
|
new_block.save(os.path.join(output_floder_normal, block_filename))
|
|||
|
|
|
|||
|
|
print("裁剪完成，图像块已保存。")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_images(self, path):
|
|||
|
|
imglist = []
|
|||
|
|
lines = os.listdir(path)
|
|||
|
|
# print(lines)
|
|||
|
|
for ii, line in enumerate(lines):
|
|||
|
|
name = line
|
|||
|
|
_imagepath = os.path.join(path, name)
|
|||
|
|
assert os.path.isfile(_imagepath)
|
|||
|
|
image = Image.open(_imagepath)
|
|||
|
|
orininal_h = image.size[1]
|
|||
|
|
orininal_w = image.size[0]
|
|||
|
|
item = {"name": name, "orininal_h": orininal_h, "orininal_w": orininal_w, "image": image}
|
|||
|
|
imglist.append(item)
|
|||
|
|
# print("共监测到{}张原始图像和标签".format(len(imglist)))
|
|||
|
|
return imglist
|
|||
|
|
|
|||
|
|
def folder_segment(self,path,net):
|
|||
|
|
# file_name = os.path.basename(path)
|
|||
|
|
# file_directory = os.path.dirname(path)
|
|||
|
|
# ori_path = os.path.join(file_directory,'ori')
|
|||
|
|
imglist = self.get_images(path)
|
|||
|
|
assert len(imglist) != 0
|
|||
|
|
for i in tqdm(imglist):
|
|||
|
|
image = i["image"]
|
|||
|
|
name = i["name"]
|
|||
|
|
orininal_w = i["orininal_w"]
|
|||
|
|
orininal_h = i["orininal_h"]
|
|||
|
|
old_img = copy.deepcopy(image)
|
|||
|
|
imaged = cv.resize(np.array(image), dsize=(512, 512), interpolation=cv.INTER_LINEAR)
|
|||
|
|
image_data = np.expand_dims(
|
|||
|
|
np.transpose(self.preprocess_input(np.array(imaged, np.float32), md=False), (2, 0, 1)), 0)
|
|||
|
|
with torch.no_grad():
|
|||
|
|
images = torch.from_numpy(image_data)
|
|||
|
|
image = images.to(device="cuda", dtype=torch.float32)
|
|||
|
|
model = net.to(device="cuda")
|
|||
|
|
out = model(image) # batch_size, 2, 512, 512
|
|||
|
|
if isinstance(out, list) or isinstance(out, tuple): # 可能有多个输出（这里把辅助解码头的也输出的所以是多个）
|
|||
|
|
out = out[0] # 就取第一个
|
|||
|
|
out = out[0] # 去掉batch
|
|||
|
|
|
|||
|
|
pr = F.softmax(out.permute(1, 2, 0), dim=-1).cpu().numpy()
|
|||
|
|
result = pr.argmax(axis=-1)
|
|||
|
|
# 结果图
|
|||
|
|
path_ = Path(path)
|
|||
|
|
parent_directory = path_.parent
|
|||
|
|
directory = os.path.join(parent_directory,'binary')
|
|||
|
|
if not os.path.exists(directory):
|
|||
|
|
os.makedirs(directory)
|
|||
|
|
output_im = Image.fromarray(np.uint8(result)).convert('P')
|
|||
|
|
|
|||
|
|
output_im.putpalette(self.palette)
|
|||
|
|
tmp_path_binary= os.path.join(directory,name)
|
|||
|
|
output_im.save(tmp_path_binary)
|
|||
|
|
|
|||
|
|
directory = os.path.join(parent_directory,'fusion')
|
|||
|
|
if not os.path.exists(directory):
|
|||
|
|
os.makedirs(directory)
|
|||
|
|
PALETTE = [(0, 0, 0), (255, 0, 0), (0, 255, 0)] # 这里主动改分类了，将识别图像与底图融合
|
|||
|
|
seg_img0 = np.reshape(np.array(PALETTE, np.uint8)[np.reshape(result, [-1])],
|
|||
|
|
[orininal_h, orininal_w, -1])
|
|||
|
|
|
|||
|
|
image0 = Image.fromarray(np.uint8(seg_img0))
|
|||
|
|
|
|||
|
|
fusion = Image.blend(old_img, image0, 0.4) # 图像融合
|
|||
|
|
tmp_path_fusion= os.path.join(directory,name)
|
|||
|
|
fusion.save(tmp_path_fusion)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def folder_segment_all(self,path,net,mode):
|
|||
|
|
# file_name = os.path.basename(path)
|
|||
|
|
# file_directory = os.path.dirname(path)
|
|||
|
|
# ori_path = os.path.join(file_directory,'ori')
|
|||
|
|
imglist = self.get_images(path)
|
|||
|
|
assert len(imglist) != 0
|
|||
|
|
for i in tqdm(imglist):
|
|||
|
|
image = i["image"]
|
|||
|
|
name = i["name"]
|
|||
|
|
orininal_w = i["orininal_w"]
|
|||
|
|
orininal_h = i["orininal_h"]
|
|||
|
|
old_img = copy.deepcopy(image)
|
|||
|
|
imaged = cv.resize(np.array(image), dsize=(512, 512), interpolation=cv.INTER_LINEAR)
|
|||
|
|
image_data = np.expand_dims(
|
|||
|
|
np.transpose(self.preprocess_input(np.array(imaged, np.float32), md=False), (2, 0, 1)), 0)
|
|||
|
|
with torch.no_grad():
|
|||
|
|
images = torch.from_numpy(image_data)
|
|||
|
|
image = images.to(device="cuda", dtype=torch.float32)
|
|||
|
|
model = net.to(device="cuda")
|
|||
|
|
out = model(image) # batch_size, 2, 512, 512
|
|||
|
|
if isinstance(out, list) or isinstance(out, tuple): # 可能有多个输出（这里把辅助解码头的也输出的所以是多个）
|
|||
|
|
out = out[0] # 就取第一个
|
|||
|
|
out = out[0] # 去掉batch
|
|||
|
|
|
|||
|
|
pr = F.softmax(out.permute(1, 2, 0), dim=-1).cpu().numpy()
|
|||
|
|
result = pr.argmax(axis=-1)
|
|||
|
|
# 结果图
|
|||
|
|
path_ = Path(path)
|
|||
|
|
parent_directory = path_.parent
|
|||
|
|
directory = os.path.join(parent_directory, mode+'_binary')
|
|||
|
|
if not os.path.exists(directory):
|
|||
|
|
os.makedirs(directory)
|
|||
|
|
output_im = Image.fromarray(np.uint8(result)).convert('P')
|
|||
|
|
|
|||
|
|
output_im.putpalette(self.palette)
|
|||
|
|
tmp_path_binary= os.path.join(directory,name)
|
|||
|
|
output_im.save(tmp_path_binary)
|
|||
|
|
|
|||
|
|
# directory = os.path.join(parent_directory,mode + '_fusion')
|
|||
|
|
# if not os.path.exists(directory):
|
|||
|
|
# os.makedirs(directory)
|
|||
|
|
# PALETTE = [(0, 0, 0), (255, 0, 0), (0, 255, 0)] # 这里主动改分类了，将识别图像与底图融合
|
|||
|
|
# seg_img0 = np.reshape(np.array(PALETTE, np.uint8)[np.reshape(result, [-1])],
|
|||
|
|
# [orininal_h, orininal_w, -1])
|
|||
|
|
|
|||
|
|
# image0 = Image.fromarray(np.uint8(seg_img0))
|
|||
|
|
|
|||
|
|
# fusion = Image.blend(old_img, image0, 0.4) # 图像融合
|
|||
|
|
# tmp_path_fusion= os.path.join(directory,name)
|
|||
|
|
# fusion.save(tmp_path_fusion)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def merge_pic(self,path):
|
|||
|
|
image = Image.open(path)
|
|||
|
|
width, height = image.size
|
|||
|
|
file_name = os.path.basename(path)
|
|||
|
|
file_directory = os.path.dirname(path)
|
|||
|
|
output_image_path = os.path.join(file_directory,file_name)
|
|||
|
|
for name in ['fusion','binary']:
|
|||
|
|
file_name_ = "merge_" +name+ "_" + file_name
|
|||
|
|
input_folder = os.path.join(file_directory,name)
|
|||
|
|
# 获取所有小图片文件名
|
|||
|
|
image_files = [f for f in os.listdir(input_folder) if f.endswith('.png')]
|
|||
|
|
image_files.sort() # 按文件名排序，确保按顺序合并
|
|||
|
|
# 假设小图的尺寸是512x512
|
|||
|
|
block_size = 512
|
|||
|
|
# 计算需要的块数
|
|||
|
|
num_blocks_x = (width + block_size - 1) // block_size # 向上取整
|
|||
|
|
num_blocks_y = (height + block_size - 1) // block_size # 向上取整
|
|||
|
|
|
|||
|
|
# 创建一个新的空白图像，用于合并块
|
|||
|
|
merged_image = Image.new('RGB', (width, height))
|
|||
|
|
|
|||
|
|
# 遍历所有块并将它们粘贴到合并图像中
|
|||
|
|
for i in range(num_blocks_x):
|
|||
|
|
for j in range(num_blocks_y):
|
|||
|
|
# 计算块的文件名
|
|||
|
|
block_filename = f'block_{i}_{j}.png'
|
|||
|
|
block_path = os.path.join(input_folder, block_filename)
|
|||
|
|
|
|||
|
|
# 打开块图像
|
|||
|
|
if os.path.exists(block_path):
|
|||
|
|
block = Image.open(block_path)
|
|||
|
|
# 计算粘贴位置
|
|||
|
|
left = i * block_size
|
|||
|
|
upper = j * block_size
|
|||
|
|
merged_image.paste(block, (left, upper))
|
|||
|
|
# 保存合并后的图像
|
|||
|
|
target = os.path.join(file_directory,file_name_)
|
|||
|
|
merged_image.save(target) # 替换为你想保存的路径
|
|||
|
|
print("合并完成，保存为:", target)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def merge_pic_all(self,path):
|
|||
|
|
image = Image.open(path)
|
|||
|
|
width, height = image.size
|
|||
|
|
file_name = os.path.basename(path)
|
|||
|
|
file_directory = os.path.dirname(path)
|
|||
|
|
path_list = []
|
|||
|
|
for name in ['pv_binary','roof_binary']:
|
|||
|
|
file_name_ = "merge_" +name+ "_" + file_name
|
|||
|
|
input_folder = os.path.join(file_directory,name)
|
|||
|
|
# 获取所有小图片文件名
|
|||
|
|
image_files = [f for f in os.listdir(input_folder) if f.endswith('.png')]
|
|||
|
|
image_files.sort() # 按文件名排序，确保按顺序合并
|
|||
|
|
# 假设小图的尺寸是512x512
|
|||
|
|
block_size = 512
|
|||
|
|
# 计算需要的块数
|
|||
|
|
num_blocks_x = (width + block_size - 1) // block_size # 向上取整
|
|||
|
|
num_blocks_y = (height + block_size - 1) // block_size # 向上取整
|
|||
|
|
|
|||
|
|
# 创建一个新的空白图像，用于合并块
|
|||
|
|
merged_image = Image.new('RGB', (width, height))
|
|||
|
|
|
|||
|
|
# 遍历所有块并将它们粘贴到合并图像中
|
|||
|
|
for i in range(num_blocks_x):
|
|||
|
|
for j in range(num_blocks_y):
|
|||
|
|
# 计算块的文件名
|
|||
|
|
block_filename = f'block_{i}_{j}.png'
|
|||
|
|
block_path = os.path.join(input_folder, block_filename)
|
|||
|
|
|
|||
|
|
# 打开块图像
|
|||
|
|
if os.path.exists(block_path):
|
|||
|
|
block = Image.open(block_path)
|
|||
|
|
# 计算粘贴位置
|
|||
|
|
left = i * block_size
|
|||
|
|
upper = j * block_size
|
|||
|
|
merged_image.paste(block, (left, upper))
|
|||
|
|
# 保存合并后的图像
|
|||
|
|
target = os.path.join(file_directory,file_name_)
|
|||
|
|
merged_image.save(target) # 替换为你想保存的路径
|
|||
|
|
print("合并完成，保存为:", target)
|
|||
|
|
path_list.append(target)
|
|||
|
|
return path_list
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
def preprocess_input(self, image, md=False):
|
|||
|
|
mean = (0.231, 0.217, 0.22) # 针对北京的数据集
|
|||
|
|
std = (0.104, 0.086, 0.085)
|
|||
|
|
if md:
|
|||
|
|
image /= 255.0
|
|||
|
|
image -= mean
|
|||
|
|
image /= std
|
|||
|
|
return image
|
|||
|
|
else:
|
|||
|
|
image /= 255.0
|
|||
|
|
return image
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 将屋顶与光伏的二值图进行合并
|
|||
|
|
def merge_binary(self, path_list):
|
|||
|
|
|
|||
|
|
# 打开二值图像
|
|||
|
|
image1_path = path_list[1] # 替换为你的屋顶图像路径
|
|||
|
|
image2_path = path_list[0] # 替换为你的光伏图像路径
|
|||
|
|
file_directory = os.path.dirname(image1_path)
|
|||
|
|
# 加载图像并确保它们是RGB模式
|
|||
|
|
image1 = Image.open(image1_path).convert('RGB')
|
|||
|
|
image2 = Image.open(image2_path).convert('RGB')
|
|||
|
|
|
|||
|
|
# 创建一个新的图像，RGB模式，初始为黑色
|
|||
|
|
output_image = Image.new('RGB', image1.size, (0, 0, 0))
|
|||
|
|
|
|||
|
|
# 获取像素数据
|
|||
|
|
pixels1 = image1.load()
|
|||
|
|
pixels2 = image2.load()
|
|||
|
|
output_pixels = output_image.load()
|
|||
|
|
|
|||
|
|
# 遍历所有像素点
|
|||
|
|
for x in range(image1.width):
|
|||
|
|
for y in range(image1.height):
|
|||
|
|
# 检查两个图像在当前像素点上的颜色
|
|||
|
|
if pixels1[x, y] == (255, 0, 0) and pixels2[x, y] == (255, 0, 0): # 红色
|
|||
|
|
output_pixels[x, y] = (255, 0, 0) # 设置为红色
|
|||
|
|
else:
|
|||
|
|
output_pixels[x, y] = (0, 0, 0) # 设置为黑色
|
|||
|
|
|
|||
|
|
# 保存输出图像
|
|||
|
|
final_path = os.path.join(file_directory, 'merge_pvroof_bianry.png')
|
|||
|
|
output_image.save(final_path) # 替换为你想保存的路径
|
|||
|
|
return final_path
|
|||
|
|
|
|||
|
|
# 将二值图与原始图合并
|
|||
|
|
def merge_final(self,path_list):
|
|||
|
|
ori_path = path_list[1] # 原始图像
|
|||
|
|
binary_path = path_list[0] # 二值图像
|
|||
|
|
# 加载图像并确保它们是RGBA模式（如果需要透明度）
|
|||
|
|
ori_image = Image.open(ori_path).convert('RGBA')
|
|||
|
|
binary_image = Image.open(binary_path).convert('RGBA')
|
|||
|
|
|
|||
|
|
# 确保两张图像的大小相同
|
|||
|
|
if ori_image.size != binary_image.size:
|
|||
|
|
raise ValueError("两张图像的尺寸必须相同！")
|
|||
|
|
|
|||
|
|
# 设置混合比例，0.5表示两张图像各占50%
|
|||
|
|
alpha = 0.4
|
|||
|
|
|
|||
|
|
# 融合图像
|
|||
|
|
blended_image = Image.blend(ori_image, binary_image, alpha)
|
|||
|
|
file_directory = os.path.dirname(ori_path)
|
|||
|
|
file_name = os.path.basename(ori_path)
|
|||
|
|
final_path = os.path.join(file_directory,'fusion_'+file_name)
|
|||
|
|
|
|||
|
|
# 保存融合后的图像
|
|||
|
|
blended_image.save(final_path) # 替换为你想保存的路径
|