Spaces:
Runtime error
Runtime error
File size: 1,983 Bytes
7d9724b ac8f703 7d9724b ac8f703 7d9724b b0f0ca6 199dc75 6f4a3d8 199dc75 7d9724b 6f4a3d8 7d9724b ac8f703 7d9724b 199dc75 7d9724b ac8f703 7d9724b a499985 7d9724b |
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 |
try:
import detectron2
except:
import os
os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
from matplotlib.pyplot import axis
import requests
import numpy as np
from torch import nn
import requests
import torch
import detectron2
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
import streamlit as st
from detectron2.utils.visualizer import ColorMode
import os
import cv2
from PIL import Image, ImageOps
import numpy as np
model_path = "model_final.pth"
cfg=get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_1x.yaml")) # load the default configuration
cfg.MODEL.WEIGHTS = 'model_final.pth'
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.8
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
st.write("""
# Gun Detection
"""
)
file = st.file_uploader("Please upload an image file(JPG/PNG/JPEG format)", type=["jpg", "png","jpeg"])
st.set_option('deprecation.showfileUploaderEncoding', False)
dataset_name="guns"
classes=['guns','Gun']
MetadataCatalog.get(dataset_name).set(thing_classes=classes)
if not torch.cuda.is_available():
cfg.MODEL.DEVICE='cpu'
predictor = DefaultPredictor(cfg)
def inference(image):
#height = image.height
#img = np.array(image.resize((800, height)))
img = np.array(image)
outputs = predictor(img)
v=Visualizer(img,MetadataCatalog.get(dataset_name),scale=1)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
return out.get_image()
if file is None:
st.text("Please upload an image file")
else:
image = Image.open(file).convert('RGB')
st.image(image,use_column_width=True)
st.write("""
# Output!!
"""
)
predictions = inference(image)
st.image(predictions,use_column_width=True) |