44 lines
1.6 KiB
Python
44 lines
1.6 KiB
Python
import gradio as gr
|
||
from transformers import MaskFormerImageProcessor, MaskFormerForInstanceSegmentation
|
||
|
||
def inference(img):
|
||
# load MaskFormer fine-tuned on COCO panoptic segmentation
|
||
model_path = "maskformer-swin-large-coco"
|
||
processor = MaskFormerImageProcessor.from_pretrained(model_path)
|
||
model = MaskFormerForInstanceSegmentation.from_pretrained(model_path)
|
||
|
||
inputs = processor(images=img, return_tensors="pt")
|
||
|
||
outputs = model(**inputs)
|
||
# model predicts class_queries_logits of shape `(batch_size, num_queries)`
|
||
# and masks_queries_logits of shape `(batch_size, num_queries, height, width)`
|
||
class_queries_logits = outputs.class_queries_logits
|
||
masks_queries_logits = outputs.masks_queries_logits
|
||
|
||
# you can pass them to processor for postprocessing
|
||
result = processor.post_process_panoptic_segmentation(outputs, target_sizes=[img.size[::-1]])[0]
|
||
# we refer to the demo notebooks for visualization (see "Resources" section in the MaskFormer docs)
|
||
predicted_panoptic_map = result["segmentation"]
|
||
return predicted_panoptic_map
|
||
|
||
examples=[['example_cat.jpg']]
|
||
|
||
with gr.Blocks() as demo:
|
||
gr.Markdown(
|
||
"""
|
||
# Panorama segmentation:maskformer-swin-large-coco
|
||
这是maskformer-swin-large-coco的Gradio Demo，用于全景分割。上传你想要的图像或者点击下面的示例来加载它。
|
||
""")
|
||
with gr.Row():
|
||
image_input = gr.Image(type="pil")
|
||
text_output = gr.Textbox()
|
||
image_button = gr.Button("上传")
|
||
image_button.click(inference, inputs=image_input, outputs=text_output)
|
||
gr.Examples(examples,inputs=image_input)
|
||
|
||
demo.launch(server_name="0.0.0.0")
|
||
|
||
|
||
|
||
|