Spaces:
Runtime error
Runtime error
| import os | |
| os.system('pip install git+https://github.com/facebookresearch/detectron2.git') | |
| import gradio as gr | |
| import logging | |
| from detectron2.engine import DefaultPredictor | |
| import cv2 | |
| from detectron2.config import get_cfg | |
| from utils import add_bboxes | |
| config_file="config.yaml" | |
| cfg = get_cfg() | |
| cfg.merge_from_file(config_file) | |
| cfg.MODEL.DEVICE="cpu" | |
| cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 | |
| # cfg.MODEL.WEIGHTS = "checkpoints_model_final_imagenet_40k_synthetic.pth" | |
| def predict( | |
| model, | |
| img | |
| ): | |
| if model=="40k synthetic": | |
| weights = "checkpoints_model_final_imagenet_40k_synthetic.pth" | |
| elif model == "100k synthetic": | |
| weights = "checkpoints_model_final_imagenet_100k_synthetic.pth" | |
| else: | |
| weights = "checkpoints_model_final_imagenet_5k_synthetic.pth" | |
| cfg.MODEL.WEIGHTS=weights | |
| predictor = DefaultPredictor(cfg) | |
| im = cv2.imread(img.name) | |
| output = predictor(im) | |
| img = add_bboxes(im, output['instances'].pred_boxes, scores=output['instances'].scores) | |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| return img | |
| title = "Indoor Pet Detection" | |
| description = "This is an application trained with synthetic data from Unity Computer Vision. We trained a single class object detection model to recognize dogs using images of pets randomly posed and placed in our home interior environment. There are multiple pre-trained models trained with 5k, 40k and 100k synthetic data that you can choose for inference. <p> For more information please refer - <a target='_blank' href='https://github.com/Unity-Technologies/Indoor-Pet-Detection'> Github Project </a> </p>" | |
| examples = [ | |
| ["5k synthetic", 'example.jpg'], | |
| ["40k synthetic", 'example.jpg'], | |
| ["100k synthetic", 'example.jpg'], | |
| ["5k synthetic", 'example-2.jpg'], | |
| ["40k synthetic", 'example-2.jpg'], | |
| ["100k synthetic", 'example-2.jpg'] | |
| ] | |
| gr.Interface(predict, [gr.inputs.Dropdown(["5k synthetic", "40k synthetic", "100k synthetic"]), gr.inputs.Image(type="file")], outputs=gr.outputs.Image(type="pil"),enable_queue=True, title=title, | |
| description=description, | |
| # article=article, | |
| examples=examples).launch(debug=True) | |