import gradio as gr import spaces from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor, CLIPModel, \ BlipForConditionalGeneration, CLIPProcessor, BlipProcessor 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"), "Qwen/Qwen2-VL-2B-Instruct": Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", torch_dtype="auto", device_map="auto"), "Qwen/Qwen2-VL-1B-Instruct": Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-1B-Instruct", torch_dtype="auto", device_map="auto"), "Qwen/Qwen2-VL-5B-Instruct": Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-5B-Instruct", torch_dtype="auto", device_map="auto"), "openai/clip-vit-base-patch32": CLIPModel.from_pretrained("openai/clip-vit-base-patch32"), "Salesforce/blip-image-captioning-base": BlipForConditionalGeneration.from_pretrained( "Salesforce/blip-image-captioning-base"), } processors = { "Qwen/Qwen2-VL-7B-Instruct": AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct"), "Qwen/Qwen2-VL-2B-Instruct": AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct"), "Qwen/Qwen2-VL-1B-Instruct": AutoProcessor.from_pretrained("Qwen/Qwen2-VL-1B-Instruct"), "Qwen/Qwen2-VL-5B-Instruct": AutoProcessor.from_pretrained("Qwen/Qwen2-VL-5B-Instruct"), "openai/clip-vit-base-patch32": CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32"), "Salesforce/blip-image-captioning-base": BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base"), } 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": 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( """ # Multi-Model Object Detection Demo This demo uses various state-of-the-art models for object detection and image-text alignment tasks. **Available Models**: - **Qwen2-VL (7B, 2B, 5B, 1B)**: Vision-language models optimized for various tasks. - **BLIP**: Image captioning and visual question answering. - **CLIP**: Contrastive learning for image-text matching. - **Flamingo**: Few-shot learning for various visual tasks. - **LLaVA**: Balanced performance in visual understanding and interactive AI tasks. **Usage**: Input an image and a description of the target object you want to detect. """ ) with gr.Tab(label="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-2B-Instruct") system_prompt = gr.Textbox(label="System Prompt", value=default_system_prompt) text_input = gr.Textbox(label="User Prompt") 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=[ ["images/2024_09_10_10_56_40.png", "solve the questions in Turkish", default_system_prompt], ["images/2024_09_10_10_58_23.png", "solve the questions in Turkish", default_system_prompt], ["images/2024_09_10_10_58_40.png", "solve the questions in Turkish", default_system_prompt], ["images/2024_09_10_11_07_31.png", "Describe the questions and write python code", default_system_prompt], ["images/IMG_3644", "Describe the image", default_system_prompt], ["images/IMG_3658", "Describe the image", default_system_prompt], ["images/IMG_4028", "Describe the image", default_system_prompt], ["images/IMG_4070", "Describe the image", default_system_prompt], ["images/comics.jpeg", "Describe the image", 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)