Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import os | |
| import json | |
| import google.generativeai as genai | |
| from PIL import Image, ImageDraw, ImageFont, ImageColor | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| api_key = os.getenv("Gemini_API_Key") | |
| genai.configure(api_key=api_key) | |
| SYSTEM_INSTRUCTION = """ | |
| Return bounding boxes as a JSON array with labels. Never return masks or code fencing. Limit to 25 objects. | |
| If an object is present multiple times, name them according to their unique characteristic. | |
| Use the format: [{"box_2d": [ymin, xmin, ymax, xmax], "label": "object_name"}] | |
| """ | |
| model = genai.GenerativeModel( | |
| model_name='gemini-2.5-flash', | |
| system_instruction=SYSTEM_INSTRUCTION | |
| ) | |
| def parse_response(text): | |
| """Cleans markdown formatting if the model includes it.""" | |
| clean = text.replace("json", "").replace("", "").strip() | |
| return clean | |
| def draw_boxes(image, boxes_json): | |
| """Draws boxes on the image using normalized coordinates (0-1000).""" | |
| img = image.copy() | |
| draw = ImageDraw.Draw(img) | |
| width, height = img.size | |
| colors = ["red", "green", "blue", "yellow", "magenta", "cyan", "orange"] | |
| try: | |
| data = json.loads(boxes_json) | |
| for i, item in enumerate(data): | |
| color = colors[i % len(colors)] | |
| # Gemini order: [ymin, xmin, ymax, xmax] | |
| ymin, xmin, ymax, xmax = item["box_2d"] | |
| left = (xmin / 1000) * width | |
| top = (ymin / 1000) * height | |
| right = (xmax / 1000) * width | |
| bottom = (ymax / 1000) * height | |
| draw.rectangle([left, top, right, bottom], outline=color, width=5) | |
| draw.text((left + 5, top + 5), item.get("label", "Object"), fill=color) | |
| except Exception as e: | |
| print(f"Error drawing boxes: {e}") | |
| return img | |
| def process_image(input_img, user_instruction): | |
| if input_img is None: | |
| return None | |
| if not user_instruction.strip(): | |
| user_instruction = "Identify and label the objects in the image." | |
| # Combine user instruction with JSON formatting requirement | |
| full_prompt = f"{user_instruction} Return ONLY the JSON array of bounding boxes." | |
| response = model.generate_content([full_prompt, input_img]) | |
| json_data = parse_response(response.text) | |
| output_img = draw_boxes(input_img, json_data) | |
| return output_img | |
| demo = gr.Interface( | |
| fn=process_image, | |
| inputs=[ | |
| gr.Image(type="pil", label="Upload Image"), | |
| gr.Textbox( | |
| label="What should I find?", | |
| placeholder="e.g., 'Find all the red apples' or 'Detect the cars and pedestrians'", | |
| lines=2 | |
| ) | |
| ], | |
| outputs=gr.Image(type="pil", label="Image with Bounding Boxes"), | |
| title="Your Own Object Detector", | |
| description="Upload an image and give specific instructions on what to detect." | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |