Update app.py

app.py

@@ -2,7 +2,8 @@
 building-segmentation
 Proof of concept showing effectiveness of a fine tuned instance segmentation model for deteting buildings.
 """
-
+import os 
+os.system("pip install 'git+https://github.com/facebookresearch/detectron2.git'")
 from transformers import DetrFeatureExtractor, DetrForSegmentation
 from PIL import Image
 import gradio as gr
@@ -11,36 +12,46 @@ import torch
 import torchvision
 import detectron2
 
+# import some common detectron2 utilities
 import itertools
 import seaborn as sns
+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, DatasetCatalog
+from PIL import Image
 
 cfg = get_cfg()
+cfg.MODEL.DEVICE='cpu'
+cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 
+cfg.MODEL.WEIGHTS = "model_weights/chatswood_buildings_poc.pth"  
 
-def segment_buildings(input_image, confidence):
-
-    cfg.MODEL.WEIGHTS = "model_weights/chatswood_buildings_poc.pth"  
-    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set a custom testing threshold
-    predictor = DefaultPredictor(cfg)
+def segment_buildings(input_image):
 
+    im = cv2.imread(input_image.name)
     outputs = predictor(im)
-
     v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
-    output = v.draw_instance_predictions(outputs["instances"].to("cpu"))
-    output_image = output.get_image()[:, :, ::-1])
-    
-    return(output_image)
+    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+    return Image.fromarray(np.uint8(out.get_image())).convert('RGB')
 
 # gradio components -inputs
-gr_image_input = gr.inputs.Image()
+gr_image_input = gr.inputs.Image(type="file")
+"""
 gr_slider_confidence = gr.inputs.Slider(0,1,.1,.7,
                                         label='Set confidence threshold % for masks')
+"""
 # gradio outputs
-gr_image_output = gr.outputs.Image() 
+gr_image_output = gr.outputs.Image(type="pil") 
+
+title = "Building Segmentation"
+description = "An instance segmentation demo for identifying boundaries of buildings in aerial images using DETR (End-to-End Object Detection) model with MaskRCNN-101 backbone"
 
 # Create user interface and launch
 gr.Interface(predict_building_mask, 
-                inputs = [gr_image_input,gr_slider_confidence],
+                inputs = gr_image_input,
                 outputs = gr_image_output,
-                title = 'Building Segmentation',
-                description = "An instance segmentation webapp using DETR (End-to-End Object Detection) model with MaskRCNN-101 backbone").launch()
+                title = title,
+                enable_queue = True,
+                description = description).launch()