File size: 3,351 Bytes
aaaf972 7ab4620 aaaf972 d717fc6 aaaf972 d717fc6 aaaf972 d717fc6 aaaf972 d717fc6 aaaf972 4a8aac3 aaaf972 |
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 |
import gradio as gr
import numpy as np
import cv2
from mmseg.apis import init_model, inference_model
import torch
def process_single_img(img_bgr, model_name):
print(type(img_bgr))
palette = [
['background', [0, 0, 0]],
['red', [255, 0, 0]]
]
palette_dict = {}
for idx, each in enumerate(palette):
palette_dict[idx] = each[1]
if model_name == 'Mask2Former':
config_file = 'CVRP_configs/CVRP_mask2former.py'
checkpoint_file = 'checkpoint/Mask2Former.pth'
elif model_name == 'KNet':
config_file = 'CVRP_configs/CVRP_knet.py'
checkpoint_file = 'checkpoint/KNet.pth'
elif model_name == 'DeepLabV3+':
config_file = 'CVRP_configs/CVRP_deeplabv3plus.py'
checkpoint_file = 'checkpoint/DeepLabV3plus.pth'
elif model_name == 'Segformer':
config_file = 'CVRP_configs/CVRP_segformer.py'
checkpoint_file = 'checkpoint/Segformer.pth'
else:
return None, None
device = 'cuda:0'
model = init_model(config_file, checkpoint_file, device=device)
result = inference_model(model, img_bgr)
pred_mask = result.pred_sem_seg.data[0].cpu().numpy()
pred_mask_bgr = np.zeros((pred_mask.shape[0], pred_mask.shape[1], 3))
for idx in palette_dict.keys():
pred_mask_bgr[np.where(pred_mask == idx)] = palette_dict[idx]
pred_mask_bgr = pred_mask_bgr.astype('uint8')
pred_viz = cv2.addWeighted(img_bgr, 1, pred_mask_bgr, 1, 0)
torch.cuda.empty_cache()
return pred_viz, pred_mask_bgr
def run_segmentation(image_input, model_select):
if model_select not in ["Mask2Former", "KNet", "DeepLabV3+", "Segformer"]:
return None, None, [("No implementa", "Error"), ("", "")]
else:
color_img, binary_img = process_single_img(image_input, model_select)
return color_img, binary_img, [("", ""), ("Segmentation Finished", "normal")]
title = """<p><h1 align="center">CVRP</h1></p>"""
# 设置SAM参数
with gr.Blocks() as iface:
gr.Markdown(title)
with gr.Row():
with gr.Column():
image_input = gr.Image(interactive=True, visible=True, label="Input Image", height=360)
with gr.Row():
model_select = gr.Dropdown(choices=["Mask2Former", "KNet", "DeepLabV3+", "Segformer"], value="Mask2Former", label="Select model", visible=True)
run_button = gr.Button(value="Run", interactive=True, visible=True)
with gr.Row():
gr.Examples(
examples=[['assets/T42_1220.jpg', 'Mask2Former'], ['assets/T99_799.jpg', 'Mask2Former'], ['assets/T92_10336.jpg', 'Mask2Former']],
inputs=[image_input, model_select])
with gr.Column():
color_output = gr.Image(interactive=False, visible=True, label="Color Image", height=360)
binary_output = gr.Image(interactive=False, visible=True, label="Binary Image", height=360)
run_status = gr.HighlightedText(
value=[("Text", "Error"), ("to be", "Label 2"), ("highlighted", "Label 3")], visible=True)
run_button.click(
fn=run_segmentation,
inputs=[image_input, model_select],
outputs=[color_output, binary_output, run_status]
)
iface.launch(debug=True, server_port=6006, server_name="127.0.0.1")
|