|
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>""" |
|
|
|
|
|
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") |
|
|
