Update app.py

app.py

@@ -1,39 +1,71 @@
+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 detection import *
-
-def detect_on_image(x):
-    detector = Detector(model_type=x)
-    image_file = st.file_uploader("Upload An Image",type=['png','jpeg','jpg'])
-    if image_file is not None:
-        file_details = {"FileName":image_file.name,"FileType":image_file.type}
-        st.write(file_details)
-        img = Image.open(image_file)
-        st.image(img, caption='Uploaded Image.')
-        with open(image_file.name,mode = "wb") as f: 
-            f.write(image_file.getbuffer())         
-        st.success("Saved File")
-        detector.onImage(image_file.name)
-        img_ = Image.open("result.jpg")
-        st.image(img_, caption='Proccesed Image.')
-
-
-def main():
-    with st.expander("About the App"):
-        st.markdown( '<p style="font-size: 30px;"><strong>Welcome to my Instance Segmentation App!</strong></p>', unsafe_allow_html= True)
-        
-
-    option = st.selectbox(
-     'What Type of File do you want to work with?',
-     ('Images', ' '))
+from detectron2.utils.visualizer import ColorMode
+import os
+import cv2
+from PIL import Image, ImageOps
+import numpy as np
 
-    
-    if option == "Images":
-        st.title('Instance Segmentation for Images')
-        st.subheader("""
-    This takes an image as an input, and provides image with bounding box and mask as an output.
-    """)
-        detect_on_image('object_detection')
+model_path = "model_final.pth"
+
+cfg = get_cfg()
+cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.6
+cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4
+cfg.MODEL.WEIGHTS = model_path
+st.write("""
+         # Car Damage Detection
+         """
+         )
+file = st.file_uploader("Please upload an image file(JPG/PNG/JPEG format)", type=["jpg", "png","jpeg"])
 
+st.set_option('deprecation.showfileUploaderEncoding', False)
+
+car_metadata = MetadataCatalog.get("test1")
+car_metadata.thing_classes = ['Damage-car','Damage','Others','Undamage']
+
+if not torch.cuda.is_available():
+    cfg.MODEL.DEVICE='cpu'
+    
+predictor = DefaultPredictor(cfg)
+def inference(image):
+    #height = image.height
 
-if __name__ == '__main__':
-	main()
+    #img = np.array(image.resize((800, height)))
+    img = np.array(image)
+    outputs = predictor(img) 
+    v = Visualizer(img[:, :, ::-1],
+                   metadata=car_metadata, 
+                   scale=0.5, 
+                   instance_mode=ColorMode.IMAGE_BW  
+    )
+    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)