Upload app.py

app.py

@@ -0,0 +1,87 @@
+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/02604.jpg', 'Mask2Former'], ['assets/T92_323.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")