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 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 import uuid 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 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 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(device).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 run_example(task_prompt, image, text_input=None): prompt = task_prompt if text_input is None else task_prompt + text_input inputs = processor(text=prompt, images=image, return_tensors="pt").to(device) with torch.inference_mode(): generated_ids = model.generate(**inputs, 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] return processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.size[0], image.size[1])) def fig_to_pil(fig): buf = io.BytesIO() fig.savefig(buf, format='png', dpi=300, bbox_inches='tight') buf.seek(0) return Image.open(buf) def plot_bbox_img(image, data): fig, ax = plt.subplots(figsize=(10, 10)) ax.imshow(image) if 'bboxes' in data and 'labels' in data: bboxes, labels = data['bboxes'], data['labels'] elif 'bboxes' in data and 'bboxes_labels' in data: bboxes, labels = data['bboxes'], data['bboxes_labels'] else: return fig_to_pil(fig) for bbox, label in zip(bboxes, labels): x1, y1, x2, y2 = bbox rect = patches.Rectangle((x1, y1), x2-x1, y2-y1, linewidth=2, edgecolor='indigo', facecolor='none') ax.add_patch(rect) plt.text(x1, y1, label, color='white', fontsize=10, bbox=dict(facecolor='indigo', alpha=0.8)) ax.axis('off') return fig_to_pil(fig) def draw_poly_img(image, prediction, fill_mask=False): fig, ax = plt.subplots(figsize=(10, 10)) ax.imshow(image) for polygons, label in zip(prediction.get('polygons', []), prediction.get('labels', [])): color = random.choice(colormap) for polygon in polygons: if isinstance(polygon[0], (int, float)): polygon = [(polygon[i], polygon[i+1]) for i in range(0, len(polygon), 2)] poly = patches.Polygon(polygon, edgecolor=color, facecolor=color if fill_mask else 'none', alpha=0.5 if fill_mask else 1, linewidth=2) ax.add_patch(poly) if polygon: plt.text(polygon[0][0], polygon[0][1], label, color='white', fontsize=10, bbox=dict(facecolor=color, alpha=0.8)) ax.axis('off') return fig_to_pil(fig) def draw_ocr_bboxes(image, prediction): fig, ax = plt.subplots(figsize=(10, 10)) ax.imshow(image) bboxes, labels = prediction['quad_boxes'], prediction['labels'] for box, label in zip(bboxes, labels): color = random.choice(colormap) box_array = np.array(box).reshape(-1, 2) # respect format polygon = patches.Polygon(box_array, edgecolor=color, fill=False, linewidth=2) ax.add_patch(polygon) plt.text(box_array[0, 0], box_array[0, 1], label, color='white', fontsize=10, bbox=dict(facecolor=color, alpha=0.8)) ax.axis('off') return fig_to_pil(fig) def plot_bbox(image, data): img_draw = image.copy() draw = ImageDraw.Draw(img_draw) for bbox, label in zip(data['bboxes'], data['labels']): x1, y1, x2, y2 = bbox draw.rectangle([x1, y1, x2, y2], outline="red", width=2) draw.text((x1, y1), label, fill="white") return np.array(img_draw) @spaces.GPU(duration=130) #remains to be seen, increasing too much may leave people queueing for long def process_video(input_video_path, task_prompt): cap = cv2.VideoCapture(input_video_path) if not cap.isOpened(): return None, [] 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) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) result_file_name = f"{uuid.uuid4()}.mp4" fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(result_file_name, fourcc, fps, (frame_width, frame_height)) processed_frames = 0 frame_results = [] color_map = {} #consistency for chromakey possibility def get_color(label): if label not in color_map: color_map[label] = random.choice(colormap) return color_map[label] 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) try: result = run_example(task_prompt, pil_image) if task_prompt == "": processed_image = plot_bbox(pil_image, result['']) frame_results.append((processed_frames + 1, result[''])) elif task_prompt == "": processed_image = pil_image.copy() draw = ImageDraw.Draw(processed_image) for i, label in enumerate(result[''].get('labels', [])): draw.text((10, 10 + i*20), label, fill="white") processed_image = np.array(processed_image) frame_results.append((processed_frames + 1, result[''])) elif task_prompt in ["", ""]: if isinstance(result[task_prompt], dict) and 'polygons' in result[task_prompt]: processed_image = draw_vid_polygons(pil_image, result[task_prompt], get_color) else: processed_image = np.array(pil_image) frame_results.append((processed_frames + 1, result[task_prompt])) else: processed_image = np.array(pil_image) out.write(cv2.cvtColor(processed_image, cv2.COLOR_RGB2BGR)) processed_frames += 1 except Exception as e: print(f"Error processing frame {processed_frames + 1}: {str(e)}") processed_image = np.array(pil_image) out.write(cv2.cvtColor(processed_image, cv2.COLOR_RGB2BGR)) processed_frames += 1 cap.release() out.release() cv2.destroyAllWindows() if processed_frames == 0: return None, frame_results return result_file_name, frame_results def draw_vid_polygons(image, prediction, get_color): img_draw = image.copy() draw = ImageDraw.Draw(img_draw) for polygons, label in zip(prediction.get('polygons', []), prediction.get('labels', [])): color = get_color(label) for polygon in polygons: if isinstance(polygon[0], (int, float)): polygon = [(polygon[i], polygon[i+1]) for i in range(0, len(polygon), 2)] draw.polygon(polygon, outline=color, fill=color) if polygon: draw.text(polygon[0], label, fill="white") return np.array(img_draw) 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_img(image, result['']))), "Object Detection": ("", lambda result: (str(result['']), plot_bbox_img(image, result['']))), "Dense Region Caption": ("", lambda result: (str(result['']), plot_bbox_img(image, result['']))), "Region Proposal": ("", lambda result: (str(result['']), plot_bbox_img(image, result['']))), "Referring Expression Segmentation": ("", lambda result: (str(result['']), draw_poly_img(image, result[''], fill_mask=True))), "Region to Segmentation": ("", lambda result: (str(result['']), draw_poly_img(image, result[''], fill_mask=True))), "Open Vocabulary Detection": ("", lambda result: (str(result['']), plot_bbox_img(image, 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 def map_task_to_prompt(task): task_mapping = { "Object Detection": "", "Dense Region Caption": "", "Referring Expression Segmentation": "", "Region to Segmentation": "" } return task_mapping.get(task, "") def process_video_p(input_video, task, text_input): prompt = map_task_to_prompt(task) if task == "Referring Expression Segmentation" and text_input: prompt += text_input result, frame_results = process_video(input_video, prompt) if result is None: return None, "Error: Video processing failed. Check logs above for info.", str(frame_results) return result, result, str(frame_results) with gr.Blocks() 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_dropdown = gr.Dropdown( choices=["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) gr.Examples( examples=[ [ "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true", "Detailed Caption", "", ], [ "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true", "Object Detection", "", ], [ "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true", "Caption to Phrase Grounding", "A green car parked in front of a yellow building." ], [ "http://ecx.images-amazon.com/images/I/51UUzBDAMsL.jpg?download=true", "OCR", "" ] ], inputs=[input_img, task_dropdown, text_input], ) 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_dropdown = gr.Dropdown( choices=["Object Detection", "Dense Region Caption", "Referring Expression Segmentation", "Region to Segmentation"], label="Video Task", value="Object Detection" ) video_text_input = gr.Textbox(label="Text Input (for Referring Expression Segmentation)", visible=False) video_submit_btn = gr.Button(value="Process Video") with gr.Column(): output_video = gr.Video(label="Processed Video") frame_results_output = gr.Textbox(label="Frame Results") 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_dropdown.change(fn=update_text_input, inputs=task_dropdown, outputs=text_input) def update_video_text_input(task): return gr.update(visible=task == "Referring Expression Segmentation") video_task_dropdown.change(fn=update_video_text_input, inputs=video_task_dropdown, outputs=video_text_input) submit_btn.click(fn=process_image, inputs=[input_img, task_dropdown, text_input], outputs=[output_text, output_image]) video_submit_btn.click(fn=process_video_p, inputs=[input_video, video_task_dropdown, video_text_input], outputs=[output_video, output_video, frame_results_output]) demo.launch()