Aedelon/LangEfficientSAM

1import os2import warnings3 4import gradio as gr5import numpy as np6from PIL import Image7 8from lang_efficient_sam.LangEfficientSAM import LangEfficientSAM9from lang_efficient_sam.utils.draw_image import draw_image10 11warnings.filterwarnings("ignore")12 13model = LangEfficientSAM()14 15 16def predict(box_threshold, text_threshold, image_path, text_prompt):17    print("Predicting... ", box_threshold, text_threshold, image_path, text_prompt)18 19    image_pil = Image.open(image_path).convert("RGB")20 21    masks, boxes, phrases, logits = model.predict(image_pil, text_prompt, box_threshold, text_threshold)22 23    labels = [f"{phrase} {logit:.2f}" for phrase, logit in zip(phrases, logits)]24 25    image_array = np.asarray(image_pil)26    image = draw_image(image_array, masks, boxes, labels)27    image = Image.fromarray(np.uint8(image)).convert("RGB")28 29    return image30 31 32title = "LangEfficientSAM"33 34inputs = [35    gr.Slider(0, 1, value=0.3, label="Box threshold"),36    gr.Slider(0, 1, value=0.25, label="Text threshold"),37    gr.Image(type="filepath", label='Image'),38    gr.Textbox(lines=1, label="Text Prompt"),39]40 41outputs = [gr.Image(type="pil", label="Output Image")]42 43examples = [44    [45        0.20,46        0.20,47        os.path.join(os.path.dirname(__file__), "images", "living.jpg"),48        "fabric",49    ],50    [51        0.36,52        0.25,53        os.path.join(os.path.dirname(__file__), "images", "fruits.jpg"),54        "apple",55    ],56    [57        0.20,58        0.20,59        os.path.join(os.path.dirname(__file__), "images", "street.jpg"),60        "car",61    ]62]63 64demo = gr.Interface(fn=predict,65                    inputs=inputs,66                    outputs=outputs,67                    examples=examples,68                    title=title)69 70demo.launch(debug=False, share=False)71