imanhamid commited on
Commit
9b19640
·
1 Parent(s): 075f26a

Create new file

Browse files
Files changed (1) hide show
  1. app.py +57 -0
app.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from gradio.outputs import Label
2
+ from icevision.all import *
3
+ from icevision.models.checkpoint import *
4
+ import PIL
5
+ import gradio as gr
6
+ import os
7
+
8
+ # Load model
9
+ class_map = ClassMap(['selected_variant'])
10
+ backbone = faster_rcnn.backbones.resnet_fpn.resnet18(pretrained=True)
11
+ model = faster_rcnn.model(backbone=backbone, num_classes=len(class_map))
12
+ model.load_state_dict(torch.load('object_localization_full-ancestry.model.pth', map_location=torch.device('cpu')))
13
+
14
+ def predict(
15
+ model, image, detection_threshold: float = 0.5, mask_threshold: float = 0.5
16
+ ):
17
+ infer_ds = Dataset.from_images([image])
18
+
19
+ batch, samples = faster_rcnn.build_infer_batch(infer_ds)
20
+ preds = faster_rcnn.predict(
21
+ model=model,
22
+ batch=batch,
23
+ detection_threshold=detection_threshold
24
+ )
25
+ return samples[0]["img"], preds[0]
26
+
27
+ def show_preds(input_image):
28
+ img = PIL.Image.fromarray(input_image, "RGB")
29
+ pred_dict = model_type.end2end_detect(img, model=model,
30
+ class_map=class_map,
31
+ detection_threshold=0.5,
32
+ display_label=False,
33
+ display_bbox=True,
34
+ return_img=True,
35
+ font_size=16,
36
+ label_color="#FF59D6")
37
+ return pred_dict["img"]
38
+
39
+ # Populate examples in Gradio interface
40
+ examples = [
41
+ ['1.jpg'],
42
+ ['2.jpg'],
43
+ ['3.jpg']
44
+ ]
45
+
46
+ display_chkbox = gr.inputs.CheckboxGroup(["Label", "BBox"], label="Display")
47
+ detection_threshold_slider = gr.inputs.Slider(minimum=0, maximum=1, step=0.1, default=0.5, label="Detection Threshold")
48
+
49
+ outputs = gr.outputs.Image(type="pil")
50
+
51
+ gr_interface = gr.Interface(
52
+ fn=show_preds,
53
+ inputs=["image", display_chkbox, detection_threshold_slider],
54
+ outputs=outputs,
55
+ title='Selection Scan - Object Detection'
56
+ examples=examples)
57
+ gr_interface.launch(inline=False, share=False, debug=True)