Spaces:
Runtime error
Runtime error
File size: 2,659 Bytes
5e10c24 |
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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import os
try:
import detectron2
except:
os.system('pip install lib/detectron2')
import numpy as np
import os, json, cv2, random
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
MODEL_YAML='COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml'
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(MODEL_YAML))
#cfg.DEVICE = 'cpu'
cfg.MODEL.DEVICE = 'cpu'
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.MODEL.WEIGHTS = "weights/model_final_2d9806.pkl"
predictor = DefaultPredictor(cfg)
import gradio as gr
from PIL import Image
def infer(input_filename):
# Predictor takes BGR.
cv2_image = cv2.imread(input_filename)
v = Visualizer(cv2_image[:, :, ::-1], # Suppose RGB
MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
scale=1.2)
results = predictor(cv2_image)
output_image = v.draw_instance_predictions(results["instances"].to("cpu")).get_image()
return Image.fromarray(np.uint8(output_image)).convert('RGB')
with gr.Blocks(title="Detectron2 Object Detection - ClassCat",
css=".gradio-container {background:lightyellow;}"
) as demo:
#sample_index = gr.State([])
gr.HTML("""<div style="font-family:'Times New Roman', 'Serif'; font-size:16pt; font-weight:bold; text-align:center; color:royalblue;">Detectron2 Object Detection</div>""")
gr.HTML("""<h4 style="color:navy;">1-a. Select an example by clicking a thumbnail below.</h4>""")
gr.HTML("""<h4 style="color:navy;">1-b. Or upload an image by clicking on the canvas.</h4>""")
with gr.Row():
input_image = gr.Image(label="Input image", type="filepath")
output_image = gr.Image(label="Output image with predicted instances", type="numpy")
gr.Examples(['samples/detectron2.png', 'samples/cat.jpg', 'samples/hotdog.jpg'], inputs=input_image)
gr.HTML("""<br/>""")
gr.HTML("""<h4 style="color:navy;">2. Then, click "Infer" button to predict object instances. It will take about 15-20 seconds (on cpu)</h4>""")
send_btn = gr.Button("Infer")
send_btn.click(fn=infer, inputs=[input_image], outputs=[output_image])
gr.HTML("""<br/>""")
gr.HTML("""<h4 style="color:navy;">Reference</h4>""")
gr.HTML("""<ul>""")
gr.HTML("""<li><a href="https://colab.research.google.com/drive/16jcaJoc6bCFAQ96jDe2HwtXj7BMD_-m5" target="_blank">Detectron2 Tutorial</a>""")
gr.HTML("""</ul>""")
#demo.queue()
demo.launch() # debug=True)
### EOF ###
|