Spaces:
Runtime error
Runtime error
File size: 1,409 Bytes
955f4a2 74ae3e6 955f4a2 08aa404 955f4a2 08aa404 3c27500 08aa404 3c27500 08aa404 3c27500 08aa404 955f4a2 08aa404 23b1543 08aa404 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
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
# print(torch.__version__, torch.cuda.is_available())
# assert torch.__version__.startswith("1.9")
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(
img,
model="40k_synthetic"
):
if model=="40k_synthetic":
weights = "checkpoints_model_final_imagenet_40k_synthetic.pth"
else:
weights = "checkpoints_model_final_imagenet_40k_synthetic.pth"
cfg.MODEL.WEIGHTS=model
predictor = DefaultPredictor(cfg)
im = cv2.imread(img.name)
output = predictor(im)
img = add_bboxes(im, output['instances'].pred_boxes, scores=output['instances'].scores)
return img
title = "Pet Detection"
description = "Demo for Indoor Pet Detection"
examples = [['example.jpg']]
gr.Interface(predict, inputs=[gr.inputs.Image(type="file"), gr.inputs.Dropdown(["40k_synthetic", "100k_synthetic"])], outputs=gr.outputs.Image(type="pil"),enable_queue=True, title=title,
description=description,
# article=article,
examples=examples).launch()
|