import gradio as gr import spaces from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor from qwen_vl_utils import process_vision_info import torch import base64 from PIL import Image, ImageDraw from io import BytesIO import re models = { "Qwen/Qwen2-VL-7B-Instruct": Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", torch_dtype="auto", device_map="auto") } processors = { "Qwen/Qwen2-VL-7B-Instruct": AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct") } DESCRIPTION = "# Qwen2-VL Object Localization Demo" def image_to_base64(image): buffered = BytesIO() image.save(buffered, format="PNG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") return img_str def draw_bounding_boxes(image, bounding_boxes, outline_color="red", line_width=2): draw = ImageDraw.Draw(image) for box in bounding_boxes: xmin, ymin, xmax, ymax = box draw.rectangle([xmin, ymin, xmax, ymax], outline=outline_color, width=line_width) return image def rescale_bounding_boxes(bounding_boxes, original_width, original_height, scaled_width=1000, scaled_height=1000): x_scale = original_width / scaled_width y_scale = original_height / scaled_height rescaled_boxes = [] for box in bounding_boxes: xmin, ymin, xmax, ymax = box rescaled_box = [ xmin * x_scale, ymin * y_scale, xmax * x_scale, ymax * y_scale ] rescaled_boxes.append(rescaled_box) return rescaled_boxes @spaces.GPU def run_example(image, text_input, system_prompt, model_id="Qwen/Qwen2-VL-7B-Instruct"): model = models[model_id].eval() processor = processors[model_id] messages = [ { "role": "user", "content": [ {"type": "image", "image": f"data:image;base64,{image_to_base64(image)}"}, {"type": "text", "text": system_prompt}, {"type": "text", "text": f"detect {text_input}"}, ], } ] text = processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) image_inputs, video_inputs = process_vision_info(messages) inputs = processor( text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt", ) inputs = inputs.to("cuda") generated_ids = model.generate(**inputs, max_new_tokens=128) generated_ids_trimmed = [ out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids) ] output_text = processor.batch_decode( generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False ) print(output_text) pattern = r'\[\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\]' matches = re.findall(pattern, str(output_text)) parsed_boxes = [[int(num) for num in match] for match in matches] scaled_boxes = rescale_bounding_boxes(parsed_boxes, image.width, image.height) return output_text, parsed_boxes, draw_bounding_boxes(image, scaled_boxes) css = """ #output { height: 500px; overflow: auto; border: 1px solid #ccc; } """ default_system_prompt = "You are a helpfull assistant to detect objects in images. When asked to detect elements based on a description you return bounding boxes for all elements in the form of [xmin, ymin, xmax, ymax] whith the values beeing scaled to 1000 by 1000 pixels. When there are more than one result answer with a list of bounding boxes in the form of [[xmin, ymin, xmax, ymax], [xmin, ymin, xmax, ymax], ...]." with gr.Blocks(css=css) as demo: gr.Markdown(DESCRIPTION) with gr.Tab(label="Qwen2-VL Input"): with gr.Row(): with gr.Column(): input_img = gr.Image(label="Input Image", type="pil") model_selector = gr.Dropdown(choices=list(models.keys()), label="Model", value="Qwen/Qwen2-VL-7B-Instruct") system_prompt = gr.Textbox(label="System Prompt", value=default_system_prompt) text_input = gr.Textbox(label="Description of Localization Target") submit_btn = gr.Button(value="Submit") with gr.Column(): model_output_text = gr.Textbox(label="Model Output Text") parsed_boxes = gr.Textbox(label="Parsed Boxes") annotated_image = gr.Image(label="Annotated Image") gr.Examples( examples=[ ["assets/image2.jpg", "orange button", default_system_prompt], ], inputs=[input_img, text_input, system_prompt], outputs=[model_output_text, parsed_boxes, annotated_image], fn=run_example, cache_examples=True, label="Try examples" ) submit_btn.click(run_example, [input_img, text_input, system_prompt, model_selector], [model_output_text, parsed_boxes, annotated_image]) demo.launch(debug=True)