Spaces:
Runtime error
Runtime error
| 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 = "<div align='center'><img src='https://user-images.githubusercontent.com/27466624/222385101-516e551c-49f5-480d-a135-4b24ee6dc308.png' width='800''/><div>" \ | |
| "<p style='text-align: center'><a href='https://github.com/open-mmlab/mmyolo'>MMYOLO</a> 是一个开源的物体检测工具箱，提供了丰富的检测模型和数据增强方式。" \ | |
| "OpenMMLab YOLO series toolbox and benchmark. Implemented RTMDet, RTMDet-Rotated,YOLOv5, YOLOv6, YOLOv7, YOLOv8,YOLOX, PPYOLOE, etc.</p>" | |
| article = "<p style='text-align: center'><a href='https://github.com/open-mmlab/mmyolo'>MMYOLO</a></p>" \ | |
| "<p style='text-align: center'><a href='https://github.com/isLinXu'>gradio build by gatilin</a></a></p>" \ | |
| gr.Interface(fn=object_detection, inputs=inputs, outputs=output_image, | |
| examples=examples, | |
| title=title, | |
| description=description, article=article, allow_flagging=False).launch() | |