170 lines
6.3 KiB
Python
170 lines
6.3 KiB
Python
import itertools
|
||
import torch
|
||
import numpy as np
|
||
from torchvision import transforms
|
||
from PIL import Image, ImageOps
|
||
import cv2
|
||
from unet import Unet
|
||
from nets.U_ConvAutoencoder import U_ConvAutoencoder
|
||
from typing import Tuple, List
|
||
|
||
|
||
# 定义卷积自编码器
|
||
|
||
class PreCA:
|
||
device: torch.device = None
|
||
model: U_ConvAutoencoder = None
|
||
transform: transforms.Compose = None
|
||
|
||
@classmethod
|
||
def initialize_model(cls, u_ca_path: str) -> None:
|
||
# 实例化模型并加载权重
|
||
cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
cls.model = U_ConvAutoencoder().to(cls.device)
|
||
cls.model.load_state_dict(torch.load(u_ca_path, map_location=cls.device))
|
||
cls.model.eval()
|
||
# 图像预处理
|
||
cls.transform = transforms.Compose([
|
||
transforms.Resize((1728, 3392)),
|
||
transforms.ToTensor()
|
||
])
|
||
|
||
@classmethod
|
||
def load_image(cls, image: Image.Image) -> torch.Tensor:
|
||
image = image.convert("L")
|
||
image = cls.transform(image).unsqueeze(0) # 添加batch维度
|
||
return image.to(cls.device)
|
||
|
||
@staticmethod
|
||
def ca_smooth(image: Image.Image) -> Image.Image:
|
||
image_cv2 = np.array(image)
|
||
# 对图像进行闭运算
|
||
closed_image = cv2.morphologyEx(image_cv2, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3)))
|
||
# Step 1: 使用高斯模糊来平滑图像边缘
|
||
blurred = cv2.GaussianBlur(closed_image, (1, 1), 0)
|
||
th = cv2.threshold(blurred, 126, 255, cv2.THRESH_BINARY)[1]
|
||
|
||
eroded_image_pil = Image.fromarray(th)
|
||
return eroded_image_pil
|
||
|
||
@classmethod
|
||
def infer(cls, image: Image.Image) -> Image.Image:
|
||
image = cls.load_image(image)
|
||
with torch.no_grad():
|
||
output = cls.model(image)
|
||
output = output.squeeze(0).cpu() # 去除batch维度并移动到CPU
|
||
output_image = transforms.ToPILImage()(output)
|
||
output_image = output_image.resize((3384, 1710), Image.NEAREST)
|
||
return output_image
|
||
|
||
|
||
class PreUnet:
|
||
|
||
@staticmethod
|
||
def blend_images_with_colorize(image1: Image.Image, image2: Image.Image, alpha: float = 0.5) -> None:
|
||
red_image1 = ImageOps.colorize(image1.convert("L"), (0, 0, 0), (255, 0, 0))
|
||
green_image2 = ImageOps.colorize(image2.convert("L"), (0, 0, 0), (0, 255, 0))
|
||
blended_image = Image.blend(red_image1, green_image2, alpha)
|
||
blended_image.show()
|
||
|
||
@staticmethod
|
||
def calculate_metrics(pred_image: Image.Image, true_image: Image.Image, threshold: int = 1) -> Tuple[int, int, int]:
|
||
pred_gray = pred_image.convert('L')
|
||
true_gray = true_image.convert('L')
|
||
|
||
pred_binary = pred_gray.point(lambda x: 0 if x < threshold else 255)
|
||
true_binary = true_gray.point(lambda x: 0 if x < threshold else 255)
|
||
|
||
pred_array = np.array(pred_binary)
|
||
true_array = np.array(true_binary)
|
||
|
||
# Calculate TP, FP, FN
|
||
TP = np.sum((pred_array == 255) & (true_array == 255))
|
||
FP = np.sum((pred_array == 255) & (true_array == 0))
|
||
FN = np.sum((pred_array == 0) & (true_array == 255))
|
||
|
||
return TP, FP, FN
|
||
|
||
@staticmethod
|
||
def apply_mask(original_image, mask_imag):
|
||
# 打开原图和mask图片
|
||
original_image = original_image.convert("RGB")
|
||
mask_image = mask_imag.convert("RGB")
|
||
|
||
# 获取图片的像素数据
|
||
original_pixels = original_image.load()
|
||
mask_pixels = mask_image.load()
|
||
|
||
# 获取图片的尺寸
|
||
width, height = original_image.size
|
||
|
||
# 遍历每个像素
|
||
for y in range(height):
|
||
for x in range(width):
|
||
# 如果mask的像素是白色 (255, 255, 255)
|
||
if mask_pixels[x, y] == (255, 255, 255):
|
||
# 将原图中的对应像素改为绿色 (0, 255, 0)
|
||
original_pixels[x, y] = (0, 255, 0)
|
||
|
||
# 保存结果图片
|
||
return original_image
|
||
|
||
|
||
|
||
@classmethod
|
||
def main(cls, ca_path: str) -> None:
|
||
PreCA.initialize_model(ca_path)
|
||
import os
|
||
from tqdm import tqdm
|
||
ious: List[float] = []
|
||
img_names: List[str] = os.listdir(dir_origin_path)
|
||
for img_name in tqdm(img_names):
|
||
if img_name.lower().endswith(
|
||
('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
|
||
image_path = os.path.join(dir_origin_path, img_name)
|
||
image = Image.open(image_path)
|
||
|
||
r_image = unet.detect_image(image)
|
||
r_image = PreCA.infer(r_image) # 自编码器
|
||
r_image = PreCA.ca_smooth(r_image)
|
||
|
||
|
||
if is_save:
|
||
if not os.path.exists(dir_save_path):
|
||
os.makedirs(dir_save_path)
|
||
r_image.save(os.path.join(dir_save_path, img_name.split('.')[0] + '_bin.png'))
|
||
if is_get_iou:
|
||
label_path = os.path.join(dir_label_path, img_name.split('.')[0] + '_bin.png')
|
||
label = Image.open(label_path)
|
||
TP, FP, FN = cls.calculate_metrics(r_image, label)
|
||
iou = TP / (TP + FP + FN)
|
||
ious.append(iou)
|
||
print(f"当前iou{iou}")
|
||
|
||
# cls.blend_images_with_colorize(label, r_image)
|
||
|
||
if is_get_iou: print(f"平均iou{np.mean(ious)}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
name_classes: List[str] = ["background", "lane"]
|
||
dir_origin_path: str = r"E:\git\unet_seg\unet\original_data\dataset_A\test\img"
|
||
# 是否计算IOU，若为True必须填写dir_label_path（label的路径）
|
||
is_get_iou: bool = True
|
||
dir_label_path: str = r"E:\git\unet_seg\unet\original_data\dataset_A\test\Label"
|
||
# 是否保存预测后的图像，若为True必须填写dir_save_path（保存路径的路径）
|
||
is_save: bool = False
|
||
dir_save_path: str = "img_out/"
|
||
# 设置多尺度监督自编码器的权重路径
|
||
u_ca_path: str = 'weights/best_conv_autoencoder1.pth'
|
||
_defaults: dict = {
|
||
"model_path": 'model_data/best80.pth', # U-Net权重地址
|
||
"num_classes": 2, # 预测类别算上背景为2
|
||
"backbone": "vgg",
|
||
"input_shape": [1696, 864], # 图像大小
|
||
"mix_type": 1,
|
||
"cuda": True, # 是否启用cuda加速
|
||
}
|
||
unet: Unet = Unet(_defaults)
|
||
PreUnet.main(u_ca_path)
|