|
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 |
|
|
|
|
|
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, |
|
|
|
examples=examples).launch(debug=True) |
|
|
