diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ce87ac5 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +/venv +/.idea \ No newline at end of file diff --git a/app.py b/app.py new file mode 100644 index 0000000..7507c71 --- /dev/null +++ b/app.py @@ -0,0 +1,74 @@ +from typing import Tuple + +import gradio as gr +import supervision as sv +import torch +from PIL import Image + +from utils.florence import load_model, run_inference, FLORENCE_DETAILED_CAPTION_TASK, \ + FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK + +MARKDOWN = """ +# Florence-2 + SAM2 🔥 +""" + +DEVICE = torch.device("cuda") + +FLORENCE_MODEL, FLORENCE_PROCESSOR = load_model(device=DEVICE) +BOX_ANNOTATOR = sv.BoxAnnotator(color_lookup=sv.ColorLookup.INDEX) +LABEL_ANNOTATOR = sv.LabelAnnotator(color_lookup=sv.ColorLookup.INDEX) + + +def process( + image_input, +) -> Tuple[Image.Image, str]: + _, result = run_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_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 + ) + + output_image = image_input.copy() + output_image = BOX_ANNOTATOR.annotate(output_image, detections) + output_image = LABEL_ANNOTATOR.annotate(output_image, detections) + return output_image, caption + + +with gr.Blocks() as demo: + gr.Markdown(MARKDOWN) + with gr.Row(): + with gr.Column(): + image_input_component = gr.Image( + type='pil', label='Upload image') + 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') + + submit_button_component.click( + fn=process, + inputs=[image_input_component], + outputs=[ + image_output_component, + text_output_component + ] + ) + +demo.launch(debug=False, show_error=True, max_threads=1) diff --git a/requirements-local.txt b/requirements-local.txt new file mode 100644 index 0000000..b8f7583 --- /dev/null +++ b/requirements-local.txt @@ -0,0 +1,9 @@ +torch +einops +spaces +timm +transformers +samv2 +gradio +supervision +opencv-python \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..d137bab --- /dev/null +++ b/requirements.txt @@ -0,0 +1,8 @@ +einops +spaces +timm +transformers +samv2 +gradio +supervision +opencv-python \ No newline at end of file diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils/florence.py b/utils/florence.py new file mode 100644 index 0000000..19ceee0 --- /dev/null +++ b/utils/florence.py @@ -0,0 +1,55 @@ +import os +from typing import Union, Any, Tuple, Dict +from unittest.mock import patch + +import torch +from PIL import Image +from transformers import AutoModelForCausalLM, AutoProcessor +from transformers.dynamic_module_utils import get_imports + +FLORENCE_CHECKPOINT = "microsoft/Florence-2-base" +FLORENCE_DETAILED_CAPTION_TASK = '' +FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK = '' + + +def fixed_get_imports(filename: Union[str, os.PathLike]) -> list[str]: + """Work around for https://huggingface.co/microsoft/phi-1_5/discussions/72.""" + if not str(filename).endswith("/modeling_florence2.py"): + return get_imports(filename) + imports = get_imports(filename) + imports.remove("flash_attn") + return imports + + +def load_model( + device: torch.device, checkpoint: str = FLORENCE_CHECKPOINT +) -> Tuple[Any, Any]: + with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports): + model = AutoModelForCausalLM.from_pretrained( + checkpoint, trust_remote_code=True).to(device).eval() + processor = AutoProcessor.from_pretrained( + checkpoint, trust_remote_code=True) + return model, processor + + +def run_inference( + model: Any, + processor: Any, + device: torch.device, + image: Image, + task: str, + text: str = "" +) -> Tuple[str, Dict]: + prompt = task + text + inputs = processor(text=prompt, images=image, return_tensors="pt").to(device) + generated_ids = model.generate( + input_ids=inputs["input_ids"], + pixel_values=inputs["pixel_values"], + max_new_tokens=1024, + num_beams=3 + ) + generated_text = processor.batch_decode( + generated_ids, skip_special_tokens=False)[0] + response = processor.post_process_generation( + generated_text, task=task, image_size=image.size) + return generated_text, response