gun_detection / app.py
siddhantgore's picture
Update app.py
a499985
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)