File size: 1,986 Bytes
24910f2
a4c8bb5
 
24910f2
e3dc417
 
b17bb63
24910f2
 
 
 
 
 
a4c8bb5
24910f2
b17bb63
24910f2
 
b17bb63
 
 
 
 
24910f2
a4c8bb5
b17bb63
 
 
24910f2
b17bb63
24910f2
b17bb63
a4c8bb5
 
b17bb63
 
24910f2
 
b17bb63
 
a4c8bb5
 
b17bb63
24910f2
b17bb63
 
 
 
24910f2
 
 
b17bb63
24910f2
b17bb63
 
 
24910f2
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
"""
building-segmentation
Proof of concept showing effectiveness of a fine tuned instance segmentation model for deteting buildings.
"""
import os
import cv2
os.system("pip install 'git+https://github.com/facebookresearch/detectron2.git'")
from transformers import DetrFeatureExtractor, DetrForSegmentation
from PIL import Image
import gradio as gr
import numpy as np
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

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

    im = cv2.imread(input_image.name)
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
    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(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(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,
                outputs = gr_image_output,
                title = title,
                enable_queue = True,
                description = description).launch()