import gradio as gr from transformers import AutoProcessor, AutoModelForCausalLM import torch import spaces import os import requests import copy from PIL import Image, ImageDraw, ImageFont import io import matplotlib.pyplot as plt import matplotlib.patches as patches import random import numpy as np device = "cpu" #device = "cuda" if torch.cuda.is_available() else "cpu" #import subprocess #subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True) #workaround for unnecessary flash_attn requirement from unittest.mock import patch from transformers.dynamic_module_utils import get_imports def fixed_get_imports(filename: str | os.PathLike) -> list[str]: if not str(filename).endswith("modeling_florence2.py"): return get_imports(filename) imports = get_imports(filename) imports.remove("flash_attn") return imports models = {} processors = {} model_ids = [ 'microsoft/Florence-2-large-ft', 'microsoft/Florence-2-large', 'microsoft/Florence-2-base-ft', 'microsoft/Florence-2-base' ] with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports): for mid in model_ids: processors[mid] = AutoProcessor.from_pretrained(mid, trust_remote_code=True) models[mid] = AutoModelForCausalLM.from_pretrained( mid, attn_implementation="sdpa", trust_remote_code=True ).to(device).eval() DESCRIPTION = "# [Florence-2 Demo](https://huggingface.co/microsoft/Florence-2-large)" colormap = ['blue','orange','green','purple','brown','pink','gray','olive','cyan','red', 'lime','indigo','violet','aqua','magenta','coral','gold','tan','skyblue'] def fig_to_pil(fig): buf = io.BytesIO() fig.savefig(buf, format='png') buf.seek(0) return Image.open(buf) # @spaces.GPU #[uncomment to use ZeroGPU] def run_example(task_prompt, image, text_input=None, model_id='microsoft/Florence-2-large', progress=gr.Progress(track_tqdm=True)): model = models[model_id] processor = processors[model_id] if text_input is None: prompt = task_prompt else: prompt = task_prompt + text_input 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, early_stopping=False, do_sample=False, num_beams=3, ) generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0] parsed_answer = processor.post_process_generation( generated_text, task=task_prompt, image_size=(image.width, image.height) ) return parsed_answer def plot_bbox(image, data): fig, ax = plt.subplots() ax.imshow(image) for bbox, label in zip(data['bboxes'], data['labels']): x1, y1, x2, y2 = bbox rect = patches.Rectangle((x1, y1), x2-x1, y2-y1, linewidth=1, edgecolor='r', facecolor='none') ax.add_patch(rect) plt.text(x1, y1, label, color='white', fontsize=8, bbox=dict(facecolor='red', alpha=0.5)) ax.axis('off') return fig def draw_polygons(image, prediction, fill_mask=False): draw = ImageDraw.Draw(image) scale = 1 for polygons, label in zip(prediction['polygons'], prediction['labels']): color = random.choice(colormap) fill_color = random.choice(colormap) if fill_mask else None for _polygon in polygons: _polygon = np.array(_polygon).reshape(-1, 2) if len(_polygon) < 3: print('Invalid polygon:', _polygon) continue _polygon = (_polygon * scale).reshape(-1).tolist() if fill_mask: draw.polygon(_polygon, outline=color, fill=fill_color) else: draw.polygon(_polygon, outline=color) draw.text((_polygon[0] + 8, _polygon[1] + 2), label, fill=color) return image def convert_to_od_format(data): bboxes = data.get('bboxes', []) labels = data.get('bboxes_labels', []) od_results = { 'bboxes': bboxes, 'labels': labels } return od_results def draw_ocr_bboxes(image, prediction): scale = 1 draw = ImageDraw.Draw(image) bboxes, labels = prediction['quad_boxes'], prediction['labels'] for box, label in zip(bboxes, labels): color = random.choice(colormap) new_box = (np.array(box) * scale).tolist() draw.polygon(new_box, width=3, outline=color) draw.text((new_box[0]+8, new_box[1]+2), "{}".format(label), align="right", fill=color) return image import json def process_image(image, task_prompt, text_input=None, model_id='microsoft/Florence-2-large'): image = Image.fromarray(image) # Convert NumPy array to PIL Image results = {} output_image = None if task_prompt == 'Caption': results = run_example('', image, model_id=model_id) elif task_prompt == 'Detailed Caption': results = run_example('', image, model_id=model_id) elif task_prompt == 'More Detailed Caption': results = run_example('', image, model_id=model_id) elif task_prompt == 'Caption + Grounding': caption = run_example('', image, model_id=model_id)[''] results = run_example('', image, caption, model_id) results[''] = caption fig = plot_bbox(image, results['']) output_image = fig_to_pil(fig) elif task_prompt == 'Detailed Caption + Grounding': caption = run_example('', image, model_id=model_id)[''] results = run_example('', image, caption, model_id) results[''] = caption fig = plot_bbox(image, results['']) output_image = fig_to_pil(fig) elif task_prompt == 'More Detailed Caption + Grounding': caption = run_example('', image, model_id=model_id)[''] results = run_example('', image, caption, model_id) results[''] = caption fig = plot_bbox(image, results['']) output_image = fig_to_pil(fig) elif task_prompt == 'Object Detection': results = run_example('', image, model_id=model_id) fig = plot_bbox(image, results['']) output_image = fig_to_pil(fig) elif task_prompt == 'Dense Region Caption': results = run_example('', image, model_id=model_id) fig = plot_bbox(image, results['']) output_image = fig_to_pil(fig) elif task_prompt == 'Region Proposal': results = run_example('', image, model_id=model_id) fig = plot_bbox(image, results['']) output_image = fig_to_pil(fig) elif task_prompt == 'Caption to Phrase Grounding': results = run_example('', image, text_input, model_id) fig = plot_bbox(image, results['']) output_image = fig_to_pil(fig) elif task_prompt == 'Referring Expression Segmentation': results = run_example('', image, text_input, model_id) output_image = draw_polygons(image.copy(), results[''], fill_mask=True) elif task_prompt == 'Region to Segmentation': results = run_example('', image, text_input, model_id) output_image = draw_polygons(image.copy(), results[''], fill_mask=True) elif task_prompt == 'Open Vocabulary Detection': results = run_example('', image, text_input, model_id) bbox_results = convert_to_od_format(results['']) fig = plot_bbox(image, bbox_results) output_image = fig_to_pil(fig) elif task_prompt == 'Region to Category': results = run_example('', image, text_input, model_id) elif task_prompt == 'Region to Description': results = run_example('', image, text_input, model_id) elif task_prompt == 'OCR': results = run_example('', image, model_id=model_id) elif task_prompt == 'OCR with Region': results = run_example('', image, model_id=model_id) output_image = draw_ocr_bboxes(image.copy(), results['']) # Default: empty result else: results = {} # ✅ Single return point return json.dumps(results), output_image css = """ #col-container { margin: 0 auto; max-width: 640px; } """ single_task_list =[ 'Caption', 'Detailed Caption', 'More Detailed Caption', 'Object Detection', 'Dense Region Caption', 'Region Proposal', 'Caption to Phrase Grounding', 'Referring Expression Segmentation', 'Region to Segmentation', 'Open Vocabulary Detection', 'Region to Category', 'Region to Description', 'OCR', 'OCR with Region' ] cascased_task_list =[ 'Caption + Grounding', 'Detailed Caption + Grounding', 'More Detailed Caption + Grounding' ] def update_task_dropdown(choice): if choice == 'Cascased task': return gr.Dropdown(choices=cascased_task_list, value='Caption + Grounding') else: return gr.Dropdown(choices=single_task_list, value='Caption') with gr.Blocks(css=css) as demo: gr.Markdown(DESCRIPTION) with gr.Tab(label="Florence-2 Image Captioning"): with gr.Row(): with gr.Column(): input_img = gr.Image(label="Input Picture") model_selector = gr.Dropdown(choices=list(models.keys()), label="Model", value='microsoft/Florence-2-large') task_type = gr.Radio(choices=['Single task', 'Cascased task'], label='Task type selector', value='Single task') task_prompt = gr.Dropdown(choices=single_task_list, label="Task Prompt", value="Caption") task_type.change(fn=update_task_dropdown, inputs=task_type, outputs=task_prompt) text_input = gr.Textbox(label="Text Input (optional)") submit_btn = gr.Button(value="Submit") with gr.Column(): output_text = gr.Textbox(label="Output Text") output_img = gr.Image(label="Output Image") gr.Examples( examples=[ ["image1.jpg", 'Object Detection'], ["image2.jpg", 'OCR with Region'] ], inputs=[input_img, task_prompt], outputs=[output_text, output_img], fn=process_image, cache_examples=True, label='Try examples' ) submit_btn.click(process_image, [input_img, task_prompt, text_input, model_selector], [output_text, output_img]) demo.launch()