Spaces:
Sleeping
Sleeping
File size: 5,572 Bytes
fa7be76 3692a8e fa7be76 3692a8e fa7be76 3692a8e fa7be76 43cb0ae fa7be76 3692a8e fa7be76 e4a7f6f fa7be76 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
# -*- coding: utf-8 -*-
# @Time : 2024/8/4 下午2:38
# @Author : xiaoshun
# @Email : [email protected]
# @File : app.py
# @Software: PyCharm
from glob import glob
import albumentations as albu
import gradio as gr
import numpy as np
import torch
from PIL import Image
from albumentations.pytorch.transforms import ToTensorV2
from src.models.components.cdnetv1 import CDnetV1
from src.models.components.cdnetv2 import CDnetV2
from src.models.components.dbnet import DBNet
from src.models.components.hrcloudnet import HRCloudNet
from src.models.components.kappamask import KappaMask
from src.models.components.mcdnet import MCDNet
from src.models.components.scnn import SCNN
from src.models.components.unetmobv2 import UNetMobV2
class Application:
def __init__(self):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.models = {
"cdnetv1": CDnetV1(num_classes=2).to(self.device),
"cdnetv2": CDnetV2(num_classes=2).to(self.device),
"hrcloudnet": HRCloudNet(num_classes=2).to(self.device),
"mcdnet": MCDNet(in_channels=3, num_classes=2).to(self.device),
"scnn": SCNN(num_classes=2).to(self.device),
"dbnet": DBNet(img_size=256, in_channels=3, num_classes=2).to(
self.device
),
"unetmobv2": UNetMobV2(num_classes=2).to(self.device),
"kappamask":KappaMask(num_classes=2,in_channels=3).to(self.device)
}
self.__load_weight()
self.transform = albu.Compose(
[
albu.Resize(256, 256, always_apply=True),
albu.ToFloat(),
ToTensorV2(),
]
)
def __load_weight(self):
"""
将模型权重加载进来
"""
for model_name, model in self.models.items():
weight_path = glob(
f"logs/train/runs/*{model_name}*/*/checkpoints/*epoch*.ckpt"
)[0]
weight = torch.load(weight_path, map_location=self.device)
state_dict = {}
for key, value in weight["state_dict"].items():
new_key = key[4:]
state_dict[new_key] = value
model.load_state_dict(state_dict)
model.eval()
print(f"{model_name} weight loaded!")
@torch.no_grad
def inference(self, image: torch.Tensor, model_name: str):
x = image.float()
x = x.unsqueeze(0)
x = x.to(self.device)
logits = self.models[model_name](x)
if isinstance(logits, tuple):
logits = logits[0]
fake_mask = torch.argmax(logits, 1).detach().cpu().squeeze(0).numpy()
return fake_mask
def give_colors_to_mask(self, mask: np.ndarray):
"""
赋予mask颜色
"""
assert len(mask.shape) == 2, "Value Error,mask的形状为(height,width)"
colors_mask = np.zeros((mask.shape[0], mask.shape[1], 3)).astype(np.float32)
colors = ((255, 255, 255), (128, 192, 128))
for color in range(2):
segc = mask == color
colors_mask[:, :, 0] += segc * (colors[color][0])
colors_mask[:, :, 1] += segc * (colors[color][1])
colors_mask[:, :, 2] += segc * (colors[color][2])
return colors_mask
def to_pil(self, image: np.ndarray, width=None, height=None):
colors_np = self.give_colors_to_mask(image)
pil_np = Image.fromarray(np.uint8(colors_np))
if width and height:
pil_np = pil_np.resize((width, height))
return pil_np
def flip(self, image_pil: Image.Image, model_name: str):
if image_pil is None:
return Image.fromarray(np.uint8(np.random.random((32, 32, 3)) * 255)), "请上传一张图片"
if model_name is None:
return Image.fromarray(np.uint8(np.random.random((32, 32, 3)) * 255)), "请选择模型名称"
image = np.array(image_pil)
raw_height, raw_width = image.shape[0], image.shape[1]
print("image type:",image.dtype)
transform = self.transform(image=image)
image = transform["image"]
fake_image = self.inference(image, model_name)
fake_image = self.to_pil(fake_image, raw_width, raw_height)
return fake_image, "success"
def tiff_to_png(image: Image.Image):
if image.format == "TIFF":
image = image.convert("RGB")
return np.array(image)
def run(self):
app = gr.Interface(
self.flip,
[
gr.Image(sources=["clipboard", "upload"], type="pil"),
gr.Radio(
["cdnetv1", "cdnetv2", "hrcloudnet", "mcdnet", "scnn", "dbnet", "unetmobv2","kappamask"],
label="model_name",
info="选择使用的模型",
),
],
[gr.Image(), gr.Textbox(label="提示信息")],
examples=[
["examples_png/barren_11.png", "dbnet"],
["examples_png/snow_10.png", "scnn"],
["examples_png/vegetation_21.png", "cdnetv2"],
["examples_png/water_22.png", "hrcloudnet"],
],
title="云检测模型在线演示",
submit_btn=gr.Button("Submit", variant="primary")
)
app.launch(share=True)
if __name__ == "__main__":
app = Application()
app.run()
|