mmdetection简单GUI检测界面

使用pysimplegui构建了简单的GUI界面，需要输入config，checkpoint和待检测的图片，指定检测设备（GPU或CPU）,指定阈值。

import io
import os
import PySimpleGUI as sg
from PIL import Image
import numpy as np
import warnings
import mmcv
import torch
warnings.filterwarnings("ignore")
from mmdet.apis import (async_inference_detector, inference_detector,
                        init_detector, show_result_pyplot)

config_file = r"D:\Github\mmdetection\work_dirs\cherry_yolox_s_8x8_300e_coco\cherry_yolox_s_8x8_300e_coco.py"
ckpt_file = r"D:\Github\mmdetection\work_dirs\cherry_yolox_s_8x8_300e_coco\epoch_10.pth"
img_file = r"D:\Datasets\cherrylabels\images\IMG_0001.jpg"

file_types = [("JPEG (*.jpg)", "*.jpg"),
              ("All files (*.*)", "*.*")]

height, width = 600,600
def main():

    layout = 	[
# 		[sg.Text('MMDetection',justification='center',pad=(200,0))],
        
		[sg.Text('config file',size=(16,1)), sg.In(config_file,size=(80,1), key='config'), sg.FileBrowse()],
        [sg.Text('model checkpoint',size=(16,1), auto_size_text=False), sg.In(ckpt_file,size=(80,1), key='ckpt'), sg.FileBrowse()],
		[sg.Text('Path to image',size=(16,1)), sg.In(img_file,size=(80,1), key='image'), sg.FileBrowse()],
        
		[sg.Text('Device',size=(16,1)), sg.Combo(['cuda:0','cpu'],default_value='cpu',key='device')],
		[sg.Text('Score threshold',size=(16,1)), sg.Slider(range=(0,1), orientation='h', resolution=0.1, default_value=0.3, size=(15,15), key='threshold')],
		[sg.OK('Detect'), sg.Exit()],
        
        [sg.Image(key="source_image",size=(height,width),background_color='grey'),
        sg.Image(key="detected_image",size=(height,width),background_color='gray')],
        ]

    window = sg.Window('MMDetection', 
                   layout,
                   default_element_size=(14,2),
                   text_justification='right',
                   auto_size_text=False,
                   # element_justification='c'
                   )

    while True:
        event, values = window.read()
        config = values['config']
        ckpt = values['ckpt']
        device = values['device']
        input_img = values["image"]
        score_thr = values['threshold']

        if os.path.exists(input_img):
                # image = Image.open(input_img)
                image = Image.open(input_img)
                image.thumbnail((height,width))
                bio = io.BytesIO()
                image.save(bio, format="PNG")
                window["source_image"].update(data=bio.getvalue())

        if event == "Exit" or event == sg.WIN_CLOSED:
            break

        if event == "Detect":

            model = init_detector(config, ckpt, device=device)
            result = inference_detector(model, input_img)
            out = 'detect.png'
            show_result_pyplot(model, input_img, result, score_thr=score_thr, out_file=out)
            
            # image = Image.fromarray(img)
            image = Image.open(out)
            image.thumbnail((height,width))
            bio = io.BytesIO()
            image.save(bio, format="PNG")
            window["detected_image"].update(data=bio.getvalue())

    window.close()


if __name__ == "__main__":

    main()

效果如下：