|
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")) |
|
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): |
|
|
|
|
|
|
|
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) |