import os os.system("pip uninstall mmcv-full") os.system("mim install 'mmengine>=0.6.0'") os.system("mim install 'mmcv>=2.0.0rc4,<2.1.0'") os.system("mim install 'mmdet>=3.0.0,<4.0.0'") os.system("mim install 'mmyolo'") import fnmatch import glob import os import PIL.Image import cv2 import gradio as gr from argparse import Namespace from pathlib import Path import mmcv import torch from mmdet.apis import inference_detector, init_detector from mmengine.config import Config, ConfigDict from mmengine.logging import print_log from mmengine.utils import ProgressBar, path from mmyolo.registry import VISUALIZERS from mmyolo.utils import switch_to_deploy from mmyolo.utils.labelme_utils import LabelmeFormat from mmyolo.utils.misc import get_file_list, show_data_classes from mim import download import warnings warnings.filterwarnings("ignore") ckpt_path = "./checkpoint" if not os.path.exists(ckpt_path): os.makedirs(ckpt_path) model_list = ['yolov5_n-v61_syncbn_fast_8xb16-300e_coco', 'yolov5_s-v61_syncbn_fast_8xb16-300e_coco', 'yolov5_m-v61_syncbn_fast_8xb16-300e_coco', 'yolov5_l-v61_syncbn_fast_8xb16-300e_coco', 'yolov5_x-v61_syncbn_fast_8xb16-300e_coco', 'yolov5_n-p6-v62_syncbn_fast_8xb16-300e_coco', 'yolov5_s-p6-v62_syncbn_fast_8xb16-300e_coco', 'yolov5_m-p6-v62_syncbn_fast_8xb16-300e_coco', 'yolov5_l-p6-v62_syncbn_fast_8xb16-300e_coco', 'yolov5_n-v61_fast_1xb64-50e_voc', 'yolov5_s-v61_fast_1xb64-50e_voc', 'yolov5_m-v61_fast_1xb64-50e_voc', 'yolov5_l-v61_fast_1xb32-50e_voc', 'yolov5_n_mask-refine-v61_syncbn_fast_8xb16-300e_coco', 'yolov5_s_mask-refine-v61_syncbn_fast_8xb16-300e_coco', 'yolov5_m_mask-refine-v61_syncbn_fast_8xb16-300e_coco', 'yolov5_l_mask-refine-v61_syncbn_fast_8xb16-300e_coco', 'yolov5_x_mask-refine-v61_syncbn_fast_8xb16-300e_coco', 'yolov5_ins_n-v61_syncbn_fast_8xb16-300e_coco_instance', 'yolov5_ins_s-v61_syncbn_fast_8xb16-300e_coco_instance', 'yolov5_ins_s-v61_syncbn_fast_non_overlap_8xb16-300e_coco_instance', 'yolov5_ins_m-v61_syncbn_fast_8xb16-300e_coco_instance', 'yolov5_ins_l-v61_syncbn_fast_8xb16-300e_coco_instance', 'yolov5_ins_x-v61_syncbn_fast_8xb16-300e_coco_instance', 'yolov6_s_syncbn_fast_8xb32-400e_coco', 'yolov6_n_syncbn_fast_8xb32-400e_coco', 'yolov6_t_syncbn_fast_8xb32-400e_coco', 'yolov6_m_syncbn_fast_8xb32-300e_coco', 'yolov6_l_syncbn_fast_8xb32-300e_coco', 'yolox_tiny_fast_8xb8-300e_coco', 'yolox_s_fast_8xb8-300e_coco', 'yolox_m_fast_8xb8-300e_coco', 'yolox_l_fast_8xb8-300e_coco', 'yolox_x_fast_8xb8-300e_coco', 'yolox_tiny_fast_8xb32-300e-rtmdet-hyp_coco', 'yolox_s_fast_8xb32-300e-rtmdet-hyp_coco', 'yolox_m_fast_8xb32-300e-rtmdet-hyp_coco', 'yolox-pose_tiny_8xb32-300e-rtmdet-hyp_coco', 'yolox-pose_s_8xb32-300e-rtmdet-hyp_coco', 'yolox-pose_m_8xb32-300e-rtmdet-hyp_coco', 'yolox-pose_l_8xb32-300e-rtmdet-hyp_coco', 'rtmdet_tiny_syncbn_fast_8xb32-300e_coco', 'kd_tiny_rtmdet_s_neck_300e_coco', 'rtmdet_s_syncbn_fast_8xb32-300e_coco', 'kd_s_rtmdet_m_neck_300e_coco', 'rtmdet_m_syncbn_fast_8xb32-300e_coco', 'kd_m_rtmdet_l_neck_300e_coco', 'rtmdet_l_syncbn_fast_8xb32-300e_coco', 'kd_l_rtmdet_x_neck_300e_coco', 'rtmdet_x_syncbn_fast_8xb32-300e_coco', 'rtmdet-r_tiny_fast_1xb8-36e_dota', 'rtmdet-r_s_fast_1xb8-36e_dota', 'rtmdet-r_m_syncbn_fast_2xb4-36e_dota', 'rtmdet-r_l_syncbn_fast_2xb4-36e_dota', 'rtmdet-r_l_syncbn_fast_2xb4-aug-100e_dota', 'yolov7_tiny_syncbn_fast_8x16b-300e_coco', 'yolov7_l_syncbn_fast_8x16b-300e_coco', 'yolov7_x_syncbn_fast_8x16b-300e_coco', 'yolov7_w-p6_syncbn_fast_8x16b-300e_coco', 'yolov7_e-p6_syncbn_fast_8x16b-300e_coco', 'ppyoloe_plus_s_fast_8xb8-80e_coco', 'ppyoloe_plus_m_fast_8xb8-80e_coco', 'ppyoloe_plus_L_fast_8xb8-80e_coco', 'ppyoloe_plus_x_fast_8xb8-80e_coco', 'yolov8_n_syncbn_fast_8xb16-500e_coco', 'yolov8_s_syncbn_fast_8xb16-500e_coco', 'yolov8_m_syncbn_fast_8xb16-500e_coco', 'yolov8_l_syncbn_fast_8xb16-500e_coco', 'yolov8_x_syncbn_fast_8xb16-500e_coco', 'yolov8_n_mask-refine_syncbn_fast_8xb16-500e_coco', 'yolov8_s_mask-refine_syncbn_fast_8xb16-500e_coco', 'yolov8_m_mask-refine_syncbn_fast_8xb16-500e_coco', 'yolov8_l_mask-refine_syncbn_fast_8xb16-500e_coco', 'yolov8_x_mask-refine_syncbn_fast_8xb16-500e_coco'] def download_test_image(): # Images torch.hub.download_url_to_file( 'https://user-images.githubusercontent.com/59380685/266264420-21575a83-4057-41cf-8a4a-b3ea6f332d79.jpg', 'bus.jpg') torch.hub.download_url_to_file( 'https://user-images.githubusercontent.com/59380685/266264536-82afdf58-6b9a-4568-b9df-551ee72cb6d9.jpg', 'dogs.jpg') torch.hub.download_url_to_file( 'https://user-images.githubusercontent.com/59380685/266264600-9d0c26ca-8ba6-45f2-b53b-4dc98460c43e.jpg', 'zidane.jpg') import shutil def clear_folder(folder_path): for filename in os.listdir(folder_path): file_path = os.path.join(folder_path, filename) try: if os.path.isfile(file_path) or os.path.islink(file_path): os.unlink(file_path) elif os.path.isdir(file_path): shutil.rmtree(file_path) except Exception as e: print(f"Failed to delete {file_path}. Reason: {e}") print(f"Clear {folder_path} successfully.") def download_cfg_checkpoint_model_name(model_name): clear_folder("./checkpoint") download(package='mmyolo', configs=[model_name], dest_root='./checkpoint') def detect_objects(args): config = args.config if isinstance(config, (str, Path)): config = Config.fromfile(config) elif not isinstance(config, Config): raise TypeError('config must be a filename or Config object, ' f'but got {type(config)}') if 'init_cfg' in config.model.backbone: config.model.backbone.init_cfg = None # build the model from a config file and a checkpoint file model = init_detector( config, args.checkpoint, device=args.device, cfg_options={}) if not args.show: path.mkdir_or_exist(args.out_dir) # init visualizer visualizer = VISUALIZERS.build(model.cfg.visualizer) visualizer.dataset_meta = model.dataset_meta # get file list files, source_type = get_file_list(args.img) # get model class name dataset_classes = model.dataset_meta.get('classes') # check class name if args.class_name is not None: for class_name in args.class_name: if class_name in dataset_classes: continue show_data_classes(dataset_classes) raise RuntimeError( 'Expected args.class_name to be one of the list, ' f'but got "{class_name}"') # start detector inference progress_bar = ProgressBar(len(files)) for file in files: result = inference_detector(model, file) img = mmcv.imread(file) img = mmcv.imconvert(img, 'bgr', 'rgb') if source_type['is_dir']: filename = os.path.relpath(file, args.img).replace('/', '_') else: filename = os.path.basename(file) out_file = None if args.show else os.path.join(args.out_dir, filename) progress_bar.update() # Get candidate predict info with score threshold pred_instances = result.pred_instances[ result.pred_instances.scores > args.score_thr] visualizer.add_datasample( filename, img, data_sample=result, draw_gt=False, show=args.show, wait_time=0, out_file=out_file, pred_score_thr=args.score_thr) def object_detection(img, model_name, out_dir, device, show, score_thr, class_name): download_cfg_checkpoint_model_name(model_name) path = "./checkpoint" config = [f for f in os.listdir(path) if fnmatch.fnmatch(f, model_name + "*.py")][0] config = path + "/" + config checkpoint = [f for f in os.listdir(path) if fnmatch.fnmatch(f, model_name + "*.pth")][0] checkpoint = path + "/" + checkpoint img_path = "input_img.jpg" img.save("input_img.jpg") args = Namespace( img=img_path, config=config, checkpoint=checkpoint, out_dir=out_dir, device=device, show=show, score_thr=score_thr, class_name=class_name, ) detect_objects(args) img_out = PIL.Image.open(os.path.join(out_dir, img_path)) return img_out inputs = [ gr.inputs.Image(type="pil", label="input"), gr.inputs.Dropdown(choices=[m for m in model_list], label='Model', default='yolov5_s-v61_syncbn_fast_8xb16-300e_coco'), gr.inputs.Textbox(default="./output", label="output"), gr.inputs.Radio(["cuda:0", "cpu"], default="cpu", label="device"), gr.inputs.Checkbox(default=False, label="show"), gr.inputs.Slider(minimum=0.1, maximum=1.0, step=0.1, default=0.3, label="score_thr"), gr.inputs.Textbox(default=None, label="class_name"), ] download_test_image() examples = [ ['bus.jpg', 'yolov5_n-v61_syncbn_fast_8xb16-300e_coco', './output', "cpu", False, 0.3, None], ['dogs.jpg', 'yolov6_s_syncbn_fast_8xb32-400e_coco', './output', "cpu", False, 0.3, None], ['zidane.jpg', 'rtmdet_tiny_syncbn_fast_8xb32-300e_coco', './output', "cpu", False, 0.3, None] ] text_output = gr.outputs.Textbox(label="输出路径") src_image = gr.outputs.Image(type="pil") output_image = gr.outputs.Image(type="pil") title = "MMYOLO detection web demo" description = "
" \ "

MMYOLO 是一个开源的物体检测工具箱，提供了丰富的检测模型和数据增强方式。" \ "OpenMMLab YOLO series toolbox and benchmark. Implemented RTMDet, RTMDet-Rotated,YOLOv5, YOLOv6, YOLOv7, YOLOv8,YOLOX, PPYOLOE, etc.

" article = "

MMYOLO

" \ "

gradio build by gatilin

" \ gr.Interface(fn=object_detection, inputs=inputs, outputs=output_image, examples=examples, title=title, description=description, article=article, allow_flagging=False).launch()