Spaces:
Sleeping
Sleeping
| 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 | |
| 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) | |