diff --git a/app.py b/app.py index d035166..1bf9d60 100644 --- a/app.py +++ b/app.py @@ -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) diff --git a/utils/florence.py b/utils/florence.py index 8543bd3..265f4e2 100644 --- a/utils/florence.py +++ b/utils/florence.py @@ -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 = '' FLORENCE_DETAILED_CAPTION_TASK = '' FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK = '' +FLORENCE_OPEN_VOCABULARY_DETECTION_TASK = '' +FLORENCE_DENSE_REGION_CAPTION_TASK = '' def fixed_get_imports(filename: Union[str, os.PathLike]) -> list[str]: diff --git a/utils/modes.py b/utils/modes.py new file mode 100644 index 0000000..fb1db4b --- /dev/null +++ b/utils/modes.py @@ -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 +] diff --git a/utils/sam.py b/utils/sam.py index 8d763b6..562cc56 100644 --- a/utils/sam.py +++ b/utils/sam.py @@ -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