nodes/rdancer-ComfyUI_Florence2SAM2
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

open vocabulary detection with Florence2 + masks with SAM2

Browse Source
This commit is contained in:
SkalskiP
2024-08-01 14:42:04 +02:00
parent d1212b2598
commit 576e22aaff
4 changed files with 140 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app.py
+108 -58
View File
@@ -1,15 +1,16 @@
from typing import Tuple, Optional
import gradio as gr
import numpy as np
import supervision as sv
import torch
from PIL import Image
from utils.florence import load_florence_model, run_florence_inference, \
FLORENCE_DETAILED_CAPTION_TASK, \
FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK
from utils.sam import load_sam_model
FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK, FLORENCE_OPEN_VOCABULARY_DETECTION_TASK
from utils.modes import INFERENCE_MODES, OPEN_VOCABULARY_DETECTION, \
CAPTION_GROUNDING_MASKS
from utils.sam import load_sam_model, run_sam_inference
MARKDOWN = """
# Florence2 + SAM2 🔥
@@ -21,94 +22,122 @@ into masks.
"""
EXAMPLES = [
"https://media.roboflow.com/notebooks/examples/dog-2.jpeg",
"https://media.roboflow.com/notebooks/examples/dog-3.jpeg",
"https://media.roboflow.com/notebooks/examples/dog-4.jpeg"
[OPEN_VOCABULARY_DETECTION, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", 'straw'],
[OPEN_VOCABULARY_DETECTION, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", 'napkin'],
[OPEN_VOCABULARY_DETECTION, "https://media.roboflow.com/notebooks/examples/dog-3.jpeg", 'tail'],
[CAPTION_GROUNDING_MASKS, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", None],
[CAPTION_GROUNDING_MASKS, "https://media.roboflow.com/notebooks/examples/dog-3.jpeg", None],
]
DEVICE = torch.device("cuda")
FLORENCE_MODEL, FLORENCE_PROCESSOR = load_florence_model(device=DEVICE)
SAM_MODEL = load_sam_model(device=DEVICE)
BOX_ANNOTATOR = sv.BoxAnnotator(color_lookup=sv.ColorLookup.INDEX)
LABEL_ANNOTATOR = sv.LabelAnnotator(
color_lookup=sv.ColorLookup.INDEX,
text_position=sv.Position.CENTER_OF_MASS,
text_color=sv.Color.BLACK,
text_color=sv.Color.from_hex("#FFFFFF"),
border_radius=5
)
MASK_ANNOTATOR = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
def process(
image_input,
) -> Tuple[Optional[Image.Image], Optional[str]]:
if image_input is None:
return None, None
_, result = run_florence_inference(
model=FLORENCE_MODEL,
processor=FLORENCE_PROCESSOR,
device=DEVICE,
image=image_input,
task=FLORENCE_DETAILED_CAPTION_TASK
)
caption = result[FLORENCE_DETAILED_CAPTION_TASK]
_, result = run_florence_inference(
model=FLORENCE_MODEL,
processor=FLORENCE_PROCESSOR,
device=DEVICE,
image=image_input,
task=FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK,
text=caption
)
detections = sv.Detections.from_lmm(
lmm=sv.LMM.FLORENCE_2,
result=result,
resolution_wh=image_input.size
)
image = np.array(image_input.convert("RGB"))
SAM_MODEL.set_image(image)
mask, score, _ = SAM_MODEL.predict(box=detections.xyxy, multimask_output=False)
# dirty fix; remove this later
if len(mask.shape) == 4:
mask = np.squeeze(mask)
detections.mask = mask.astype(bool)
output_image = image_input.copy()
def annotate_image(image, detections):
output_image = image.copy()
output_image = MASK_ANNOTATOR.annotate(output_image, detections)
output_image = BOX_ANNOTATOR.annotate(output_image, detections)
output_image = LABEL_ANNOTATOR.annotate(output_image, detections)
return output_image, caption
return output_image
def on_mode_dropdown_change(text):
return [
gr.Textbox(visible=text == OPEN_VOCABULARY_DETECTION),
gr.Textbox(visible=text == CAPTION_GROUNDING_MASKS),
]
def process(
mode_dropdown, image_input, text_input
) -> Tuple[Optional[Image.Image], Optional[str]]:
if not image_input:
return None, None
if mode_dropdown == OPEN_VOCABULARY_DETECTION:
if not text_input:
return None, None
_, result = run_florence_inference(
model=FLORENCE_MODEL,
processor=FLORENCE_PROCESSOR,
device=DEVICE,
image=image_input,
task=FLORENCE_OPEN_VOCABULARY_DETECTION_TASK,
text=text_input
)
detections = sv.Detections.from_lmm(
lmm=sv.LMM.FLORENCE_2,
result=result,
resolution_wh=image_input.size
)
detections = run_sam_inference(SAM_MODEL, image_input, detections)
return annotate_image(image_input, detections), None
if mode_dropdown == CAPTION_GROUNDING_MASKS:
_, result = run_florence_inference(
model=FLORENCE_MODEL,
processor=FLORENCE_PROCESSOR,
device=DEVICE,
image=image_input,
task=FLORENCE_DETAILED_CAPTION_TASK
)
caption = result[FLORENCE_DETAILED_CAPTION_TASK]
_, result = run_florence_inference(
model=FLORENCE_MODEL,
processor=FLORENCE_PROCESSOR,
device=DEVICE,
image=image_input,
task=FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK,
text=caption
)
detections = sv.Detections.from_lmm(
lmm=sv.LMM.FLORENCE_2,
result=result,
resolution_wh=image_input.size
)
detections = run_sam_inference(SAM_MODEL, image_input, detections)
return annotate_image(image_input, detections), caption
with gr.Blocks() as demo:
gr.Markdown(MARKDOWN)
mode_dropdown_component = gr.Dropdown(
choices=INFERENCE_MODES,
value=INFERENCE_MODES[0],
label="Mode",
info="Select a mode to use.",
interactive=True
)
with gr.Row():
with gr.Column():
image_input_component = gr.Image(
type='pil', label='Upload image')
text_input_component = gr.Textbox(
label='Text prompt')
submit_button_component = gr.Button(value='Submit', variant='primary')
with gr.Column():
image_output_component = gr.Image(type='pil', label='Image output')
text_output_component = gr.Textbox(label='Caption output')
text_output_component = gr.Textbox(label='Caption output', visible=False)
submit_button_component.click(
fn=process,
inputs=[image_input_component],
outputs=[
image_output_component,
text_output_component
]
)
with gr.Row():
gr.Examples(
fn=process,
examples=EXAMPLES,
inputs=[image_input_component],
inputs=[
mode_dropdown_component,
image_input_component,
text_input_component
],
outputs=[
image_output_component,
text_output_component
@@ -116,4 +145,25 @@ with gr.Blocks() as demo:
run_on_click=True
)
submit_button_component.click(
fn=process,
inputs=[
mode_dropdown_component,
image_input_component,
text_input_component
],
outputs=[
image_output_component,
text_output_component
]
)
mode_dropdown_component.change(
on_mode_dropdown_change,
inputs=[mode_dropdown_component],
outputs=[
text_input_component,
text_output_component
]
)
demo.launch(debug=False, show_error=True, max_threads=1)
utils/florence.py
+3
View File
@@ -8,8 +8,11 @@ from transformers import AutoModelForCausalLM, AutoProcessor
from transformers.dynamic_module_utils import get_imports
FLORENCE_CHECKPOINT = "microsoft/Florence-2-large"
FLORENCE_OBJECT_DETECTION_TASK = '<OD>'
FLORENCE_DETAILED_CAPTION_TASK = '<MORE_DETAILED_CAPTION>'
FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK = '<CAPTION_TO_PHRASE_GROUNDING>'
FLORENCE_OPEN_VOCABULARY_DETECTION_TASK = '<OPEN_VOCABULARY_DETECTION>'
FLORENCE_DENSE_REGION_CAPTION_TASK = '<DENSE_REGION_CAPTION>'
def fixed_get_imports(filename: Union[str, os.PathLike]) -> list[str]:
utils/modes.py
+7
View File
@@ -0,0 +1,7 @@
OPEN_VOCABULARY_DETECTION = "open vocabulary detection + masks"
CAPTION_GROUNDING_MASKS = "caption + grounding + masks"
INFERENCE_MODES = [
OPEN_VOCABULARY_DETECTION,
CAPTION_GROUNDING_MASKS
]
utils/sam.py
+22
View File
@@ -1,4 +1,9 @@
from typing import Any
import numpy as np
import supervision as sv
import torch
from PIL import Image
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
@@ -13,3 +18,20 @@ def load_sam_model(
) -> SAM2ImagePredictor:
model = build_sam2(config, checkpoint, device=device)
return SAM2ImagePredictor(sam_model=model)
def run_sam_inference(
model: Any,
image: Image,
detections: sv.Detections
) -> sv.Detections:
image = np.array(image.convert("RGB"))
model.set_image(image)
mask, score, _ = model.predict(box=detections.xyxy, multimask_output=False)
# dirty fix; remove this later
if len(mask.shape) == 4:
mask = np.squeeze(mask)
detections.mask = mask.astype(bool)
return detections