Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 3,587 Bytes
b577d3a 725f958 b577d3a 725f958 b577d3a 725f958 b577d3a 725f958 b577d3a 725f958 b577d3a 725f958 b577d3a 725f958 b577d3a 725f958 |
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 |
import time
import gradio as gr
import numpy as np
import supervision as sv
from PIL import Image
import torch
from transformers import SamModel, SamProcessor
from typing import Tuple
MARKDOWN = """
# EfficientSAM sv. SAM
"""
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
SAM_MODEL = SamModel.from_pretrained("facebook/sam-vit-huge").to(DEVICE)
SAM_PROCESSOR = SamProcessor.from_pretrained("facebook/sam-vit-huge")
MASK_ANNOTATOR = sv.MaskAnnotator(
color=sv.Color.red(),
color_lookup=sv.ColorLookup.INDEX)
def annotate_image(image: np.ndarray, detections: sv.Detections) -> np.ndarray:
bgr_image = image[:, :, ::-1]
annotated_bgr_image = MASK_ANNOTATOR.annotate(
scene=bgr_image, detections=detections)
return annotated_bgr_image[:, :, ::-1]
def efficient_sam_inference(
image: np.ndarray,
x_min: int,
y_min: int,
x_max: int,
y_max: int
) -> np.ndarray:
time.sleep(0.2)
return image
def sam_inference(
image: np.ndarray,
x_min: int,
y_min: int,
x_max: int,
y_max: int
) -> np.ndarray:
input_boxes = [[[x_min, y_min, x_max, y_max]]]
inputs = SAM_PROCESSOR(
Image.fromarray(image),
input_boxes=[input_boxes],
return_tensors="pt"
).to(DEVICE)
with torch.no_grad():
outputs = SAM_MODEL(**inputs)
mask = SAM_PROCESSOR.image_processor.post_process_masks(
outputs.pred_masks.cpu(),
inputs["original_sizes"].cpu(),
inputs["reshaped_input_sizes"].cpu()
)[0][0][0].numpy()
mask = mask[np.newaxis, ...]
detections = sv.Detections(xyxy=sv.mask_to_xyxy(masks=mask), mask=mask)
return annotate_image(image=image, detections=detections)
def inference(
image: np.ndarray,
x_min: int,
y_min: int,
x_max: int,
y_max: int
) -> Tuple[np.ndarray, np.ndarray]:
return (
efficient_sam_inference(image, x_min, y_min, x_max, y_max),
sam_inference(image, x_min, y_min, x_max, y_max)
)
with gr.Blocks() as demo:
gr.Markdown(MARKDOWN)
with gr.Tab(label="Box prompt"):
with gr.Row():
with gr.Column():
input_image = gr.Image()
with gr.Accordion(label="Box", open=False):
with gr.Row():
x_min_number = gr.Number(label="x_min")
y_min_number = gr.Number(label="y_min")
x_max_number = gr.Number(label="x_max")
y_max_number = gr.Number(label="y_max")
efficient_sam_output_image = gr.Image()
sam_output_image = gr.Image()
with gr.Row():
submit_button = gr.Button("Submit")
gr.Examples(
fn=inference,
examples=[
[
'https://media.roboflow.com/notebooks/examples/dog.jpeg',
69,
247,
624,
930
]
],
inputs=[input_image, x_min_number, y_min_number, x_max_number, y_max_number],
outputs=[efficient_sam_output_image, sam_output_image],
)
submit_button.click(
efficient_sam_inference,
inputs=[input_image, x_min_number, y_min_number, x_max_number, y_max_number],
outputs=efficient_sam_output_image
)
submit_button.click(
sam_inference,
inputs=[input_image, x_min_number, y_min_number, x_max_number, y_max_number],
outputs=sam_output_image
)
demo.launch(debug=False, show_error=True)
|