File size: 6,365 Bytes
8fa1f84 |
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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
import os
import gradio as gr
from demo.sam_inference import SAM_Inference
from demo.seagull_inference import Seagull
from demo.mask_utils import ImageSketcher
class Main_ui():
def __init__(self, args) -> None:
self.args = args
self.seagull = Seagull(model_path=args.model)
self.example_list = self.load_example()
self.sam = SAM_Inference()
# self.sam_predictor = get_sam_predictor()
# self.mask_generator = get_mask_generator()
def load_example(self):
examples = []
for file in sorted(os.listdir(self.args.example_path)):
examples.append([os.path.join(self.args.example_path, file)])
return examples
def load_demo(self):
with gr.Blocks() as demo:
preprocessed_img = gr.State(value=None)
binary_mask = gr.State(value=None)
with gr.Row():
gr.Markdown("""
<img src="https://github.com/chencn2020/SEAGULL/raw/main/imgs/Logo/logo.png" alt="SEAGULL" style="height: auto; width: 100%; margin-bottom: 3%;">
## 🔔 Usage
Firstly, you need to upload an image and choose the analyse types **(quality score, importance score and distortion analysis)**.
Then you can click **(points)** or pull a frame **(bbox)** on the image to indicate the region of interest (ROIs).
After that, this demo process the following steps:
> 1. SAM extracts the mask-based ROIs based on your clicked points or frame.
> 2. Based on the uploaded image and mask-based ROIs, SEAGULL analyses the quality of the ROIs.
""")
with gr.TabItem("Mask-based ROIs (Points)"):
with gr.Row():
input_image_ponit = gr.Image(type="numpy", label='Input image', height=512) # input image
output_mask_ponit = gr.Image(label='Mask-based ROI', height=512) # output binary mask
with gr.Row():
output_mask_point_on_img = gr.Image(label='Mask on image', height=512) # mask on image for better view
with gr.Column():
radio_point = gr.Radio(label='Analysis type', choices=['Quality Score', 'Importance Score', 'Distortion Analysis'], value='Quality Score')
output_text_point = gr.Textbox(label='Analysis Results')
point_seg_button = gr.Button('Analysis')
point_example = gr.Dataset(label='Examples', components=[input_image_ponit], samples=self.example_list)
with gr.TabItem("Mask-based ROIs (BBox)"):
with gr.Row():
input_image_BBOX = ImageSketcher(type="numpy", label='Input image', height=512)
output_mask_BBOX = gr.Image(label='Mask-based ROI', height=512)
with gr.Row():
output_BBOX_mask_on_img = gr.Image(label='Mask on image', height=512)
with gr.Column():
radio_BBOX = gr.Radio(label='Analysis type', choices=['Quality Score', 'Importance Score', 'Distortion Analysis'], value='Quality Score')
output_text_BBOX = gr.Textbox(label='ROI Quality Analysis')
box_seg_button = gr.Button('Generate mask and analysis')
box_analyse_button = gr.Button('Analysis')
BBOX_example = gr.Dataset(label='Examples', components=[input_image_BBOX], samples=self.example_list)
# click point
input_image_ponit.upload(
self.seagull.init_image,
[input_image_ponit],
[preprocessed_img, input_image_ponit, input_image_BBOX]
)
point_example.click(
self.seagull.init_image,
[point_example],
[preprocessed_img, input_image_ponit, input_image_BBOX]
)
# after clicking on the image
input_image_ponit.select(
self.sam.img_select_point,
[preprocessed_img],
[input_image_ponit, output_mask_ponit, output_mask_point_on_img, binary_mask]
).then(
self.seagull.seagull_predict,
[preprocessed_img, binary_mask, radio_point],
[output_text_point]
)
point_seg_button.click(
self.seagull.seagull_predict,
[preprocessed_img, binary_mask, radio_point],
[output_text_point]
)
# draw frame
input_image_BBOX.upload(
self.seagull.init_image,
[input_image_BBOX],
[preprocessed_img, input_image_ponit, input_image_BBOX]
)
BBOX_example.click(
self.seagull.init_image,
[BBOX_example],
[preprocessed_img, input_image_ponit, input_image_BBOX]
)
# after drawing a frame on the image
input_image_BBOX.select(
self.sam.gen_box_seg,
[input_image_BBOX],
[output_mask_BBOX, output_BBOX_mask_on_img, binary_mask]
)
box_seg_button.click(
self.sam.gen_box_seg,
[input_image_BBOX],
[output_mask_BBOX, output_BBOX_mask_on_img, binary_mask]
).then(
self.seagull.seagull_predict,
[preprocessed_img, binary_mask, radio_BBOX],
[output_text_BBOX]
)
box_analyse_button.click(
self.seagull.seagull_predict,
[preprocessed_img, binary_mask, radio_BBOX],
[output_text_BBOX]
)
return demo |