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)