import os from unittest.mock import patch import spaces import gradio as gr from transformers import AutoProcessor, AutoModelForCausalLM from transformers.dynamic_module_utils import get_imports import torch import requests from PIL import Image, ImageDraw import random import numpy as np import matplotlib.pyplot as plt import matplotlib.patches as patches import cv2 import io def workaround_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 with patch("transformers.dynamic_module_utils.get_imports", workaround_fixed_get_imports): model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-large-ft", trust_remote_code=True).to("cuda").eval() processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large-ft", trust_remote_code=True) 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 def run_example(task_prompt, image, text_input=None): if text_input is None: prompt = task_prompt else: prompt = task_prompt + text_input inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda") with torch.inference_mode(): 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.size[0], image.size[1]) ) 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='indigo', alpha=0.5)) ax.axis('off') return fig_to_pil(fig) def draw_polygons(image, prediction, fill_mask=False): fig, ax = plt.subplots() ax.imshow(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 _polygon.shape[0] < 3: print('Invalid polygon:', _polygon) continue _polygon = (_polygon * scale).reshape(-1).tolist() if len(_polygon) % 2 != 0: print('Invalid polygon:', _polygon) continue polygon_points = np.array(_polygon).reshape(-1, 2) if fill_mask: polygon = patches.Polygon(polygon_points, edgecolor=color, facecolor=fill_color, linewidth=2) else: polygon = patches.Polygon(polygon_points, edgecolor=color, fill=False, linewidth=2) ax.add_patch(polygon) plt.text(polygon_points[0, 0], polygon_points[0, 1], label, color='white', fontsize=8, bbox=dict(facecolor=color, alpha=0.5)) ax.axis('off') return fig_to_pil(fig) def draw_ocr_bboxes(image, prediction): fig, ax = plt.subplots() ax.imshow(image) scale = 1 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() polygon = patches.Polygon(new_box, edgecolor=color, fill=False, linewidth=3) ax.add_patch(polygon) plt.text(new_box[0], new_box[1], label, color='white', fontsize=8, bbox=dict(facecolor=color, alpha=0.5)) ax.axis('off') return fig_to_pil(fig) @spaces.GPU(duration=120) def process_video(input_video_path, task_prompt): cap = cv2.VideoCapture(input_video_path) if not cap.isOpened(): print("Error: Can't open the video file.") return frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = cap.get(cv2.CAP_PROP_FPS) fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter("output_vid.mp4", fourcc, fps, (frame_width, frame_height)) while cap.isOpened(): ret, frame = cap.read() if not ret: break frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) pil_image = Image.fromarray(frame_rgb) result = run_example(task_prompt, pil_image) if task_prompt == "": processed_image = plot_bbox(pil_image, result['']) elif task_prompt == "": processed_image = plot_bbox(pil_image, result['']) else: processed_image = pil_image processed_frame = cv2.cvtColor(np.array(processed_image), cv2.COLOR_RGB2BGR) out.write(processed_frame) cap.release() out.release() cv2.destroyAllWindows() return "output_vid.mp4" css = """ #output { min-height: 100px; overflow: auto; border: 1px solid #ccc; } """ with gr.Blocks(css=css) as demo: gr.HTML("

Microsoft Florence-2-large-ft

") with gr.Tab(label="Image"): with gr.Row(): with gr.Column(): input_img = gr.Image(label="Input Picture", type="pil") task_radio = gr.Radio( ["Caption", "Detailed Caption", "More Detailed Caption", "Caption to Phrase Grounding", "Object Detection", "Dense Region Caption", "Region Proposal", "Referring Expression Segmentation", "Region to Segmentation", "Open Vocabulary Detection", "Region to Category", "Region to Description", "OCR", "OCR with Region"], label="Task", value="Caption" ) text_input = gr.Textbox(label="Text Input (is Optional)", visible=False) submit_btn = gr.Button(value="Submit") with gr.Column(): output_text = gr.Textbox(label="Results") output_image = gr.Image(label="Image", type="pil") with gr.Tab(label="Video"): with gr.Row(): with gr.Column(): input_video = gr.Video(label="Video") video_task_radio = gr.Radio( ["Object Detection", "Dense Region Caption"], label="Video Task", value="Object Detection" ) video_submit_btn = gr.Button(value="Process Video") with gr.Column(): output_video = gr.Video(label="Video") def update_text_input(task): return gr.update(visible=task in ["Caption to Phrase Grounding", "Referring Expression Segmentation", "Region to Segmentation", "Open Vocabulary Detection", "Region to Category", "Region to Description"]) task_radio.change(fn=update_text_input, inputs=task_radio, outputs=text_input) def process_image(image, task, text): task_mapping = { "Caption": ("", lambda result: (result[''], image)), "Detailed Caption": ("", lambda result: (result[''], image)), "More Detailed Caption": ("", lambda result: (result[''], image)), "Caption to Phrase Grounding": ("", lambda result: (str(result['']), plot_bbox(image, result['']))), "Object Detection": ("", lambda result: (str(result['']), plot_bbox(image, result['']))), "Dense Region Caption": ("", lambda result: (str(result['']), plot_bbox(image, result['']))), "Region Proposal": ("", lambda result: (str(result['']), plot_bbox(image, result['']))), "Referring Expression Segmentation": ("", lambda result: (str(result['']), draw_polygons(image, result[''], fill_mask=True))), "Region to Segmentation": ("", lambda result: (str(result['']), draw_polygons(image, result[''], fill_mask=True))), "Open Vocabulary Detection": ("", lambda result: (str(convert_to_od_format(result[''])), plot_bbox(image, convert_to_od_format(result[''])))), "Region to Category": ("", lambda result: (result[''], image)), "Region to Description": ("", lambda result: (result[''], image)), "OCR": ("", lambda result: (result[''], image)), "OCR with Region": ("", lambda result: (str(result['']), draw_ocr_bboxes(image, result['']))), } if task in task_mapping: prompt, process_func = task_mapping[task] result = run_example(prompt, image, text) return process_func(result) else: return "", image submit_btn.click(fn=process_image, inputs=[input_img, task_radio, text_input], outputs=[output_text, output_image]) video_submit_btn.click(fn=process_video, inputs=[input_video, video_task_radio], outputs=output_video) demo.launch()