Update app.py

app.py

@@ -1,100 +1,101 @@
-import google.generativeai as genai
-from PIL import Image
-import re
-import cv2
-import gradio as gr
-
-def parse_bounding_box(response):
-    bounding_boxes = re.findall(r'\[(\d+,\s*\d+,\s*\d+,\s*\d+,\s*[\w\s]+)\]', response)
-
-    # Convert each group into a list of integers and labels.
-    parsed_boxes = []
-    for box in bounding_boxes:
-        parts = box.split(',')
-        numbers = list(map(int, parts[:-1]))
-        label = parts[-1].strip()
-        parsed_boxes.append((numbers, label))
-
-    # Return the list of bounding boxes with their labels.
-    return parsed_boxes
-
-# Draw bounding boxes with labels.
-def draw_bounding_boxes(image, bounding_boxes_with_labels):
-    label_colors = {}
-    if image.mode != 'RGB':
-        image = image.convert('RGB')
-
-    image = np.array(image)
-
-    for bounding_box, label in bounding_boxes_with_labels:
-        # Normalize the bounding box coordinates
-        width, height = image.shape[1], image.shape[0]
-        ymin, xmin, ymax, xmax = bounding_box
-        x1 = int(xmin / 1000 * width)
-        y1 = int(ymin / 1000 * height)
-        x2 = int(xmax / 1000 * width)
-        y2 = int(ymax / 1000 * height)
-
-        if label not in label_colors:
-            color = np.random.randint(0, 256, (3,)).tolist()
-            label_colors[label] = color
-        else:
-            color = label_colors[label]
-
-        font = cv2.FONT_HERSHEY_SIMPLEX
-        font_scale = 1
-        font_thickness = 2
-        box_thickness = 2
-        text_size = cv2.getTextSize(label, font, font_scale, font_thickness)[0]
-
-        text_bg_x1 = x1
-        text_bg_y1 = y1 - text_size[1] - 5
-        text_bg_x2 = x1 + text_size[0] + 8
-        text_bg_y2 = y1
-
-        cv2.rectangle(image, (text_bg_x1, text_bg_y1), (text_bg_x2, text_bg_y2), color, -1)
-        cv2.putText(image, label, (x1 + 2, y1 - 5), font, font_scale, (255, 255, 255), font_thickness)
-        cv2.rectangle(image, (x1, y1), (x2, y2), color, box_thickness)
-
-    image = Image.fromarray(image)
-    return image
-
-def detect_objects(api_key, prompt, input_image):
-    genai.configure(api_key=api_key)
-
-    img = Image.open(input_image)
-
-    model = genai.GenerativeModel(model_name='gemini-1.5-pro')
-
-    response = model.generate_content([
-        img,
-        (
-            f"Return bounding boxes for {prompt} in the image in the following format as"
-            " a list. \n [ymin, xmin, ymax, xmax, object_name]. "
-        ),
-    ])
-
-    result = response.text
-    result = result[result.find('-'):].strip()
-
-    bounding_box = parse_bounding_box(result)
-    output = draw_bounding_boxes(img, bounding_box)
-    
-    return output
-
-# Gradio app
-demo = gr.Interface(
-    fn=detect_objects,
-    inputs=[
-        gr.Textbox(label="Your Gemini API Key", type="password"),
-        gr.Textbox(label="Object(s) to detect", value="famous personality"),
-        gr.Image(type="filepath", label="Input Image")
-    ],
-    outputs=gr.Image(type="pil", label="Detected Image"),
-    title="Object Detection using Gemini ✨",
-    description="Detect objects in images using the Gemini.",
-    allow_flagging="never"
-)
-
-if __name__ == "__main__":
-    demo.launch(debug=True)
+import google.generativeai as genai
+from PIL import Image
+import re
+import cv2
+import numpy as np
+import gradio as gr
+
+def parse_bounding_box(response):
+    bounding_boxes = re.findall(r'\[(\d+,\s*\d+,\s*\d+,\s*\d+,\s*[\w\s]+)\]', response)
+
+    # Convert each group into a list of integers and labels.
+    parsed_boxes = []
+    for box in bounding_boxes:
+        parts = box.split(',')
+        numbers = list(map(int, parts[:-1]))
+        label = parts[-1].strip()
+        parsed_boxes.append((numbers, label))
+
+    # Return the list of bounding boxes with their labels.
+    return parsed_boxes
+
+# Draw bounding boxes with labels.
+def draw_bounding_boxes(image, bounding_boxes_with_labels):
+    label_colors = {}
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+
+    image = np.array(image)
+
+    for bounding_box, label in bounding_boxes_with_labels:
+        # Normalize the bounding box coordinates
+        width, height = image.shape[1], image.shape[0]
+        ymin, xmin, ymax, xmax = bounding_box
+        x1 = int(xmin / 1000 * width)
+        y1 = int(ymin / 1000 * height)
+        x2 = int(xmax / 1000 * width)
+        y2 = int(ymax / 1000 * height)
+
+        if label not in label_colors:
+            color = np.random.randint(0, 256, (3,)).tolist()
+            label_colors[label] = color
+        else:
+            color = label_colors[label]
+
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        font_scale = 1
+        font_thickness = 2
+        box_thickness = 2
+        text_size = cv2.getTextSize(label, font, font_scale, font_thickness)[0]
+
+        text_bg_x1 = x1
+        text_bg_y1 = y1 - text_size[1] - 5
+        text_bg_x2 = x1 + text_size[0] + 8
+        text_bg_y2 = y1
+
+        cv2.rectangle(image, (text_bg_x1, text_bg_y1), (text_bg_x2, text_bg_y2), color, -1)
+        cv2.putText(image, label, (x1 + 2, y1 - 5), font, font_scale, (255, 255, 255), font_thickness)
+        cv2.rectangle(image, (x1, y1), (x2, y2), color, box_thickness)
+
+    image = Image.fromarray(image)
+    return image
+
+def detect_objects(api_key, prompt, input_image):
+    genai.configure(api_key=api_key)
+
+    img = Image.open(input_image)
+
+    model = genai.GenerativeModel(model_name='gemini-1.5-pro')
+
+    response = model.generate_content([
+        img,
+        (
+            f"Return bounding boxes for {prompt} in the image in the following format as"
+            " a list. \n [ymin, xmin, ymax, xmax, object_name]. "
+        ),
+    ])
+
+    result = response.text
+    result = result[result.find('-'):].strip()
+
+    bounding_box = parse_bounding_box(result)
+    output = draw_bounding_boxes(img, bounding_box)
+    
+    return output
+
+# Gradio app
+demo = gr.Interface(
+    fn=detect_objects,
+    inputs=[
+        gr.Textbox(label="Your Gemini API Key", type="password"),
+        gr.Textbox(label="Object(s) to detect", value="famous personality"),
+        gr.Image(type="filepath", label="Input Image")
+    ],
+    outputs=gr.Image(type="pil", label="Detected Image"),
+    title="Object Detection using Gemini ✨",
+    description="Detect objects in images using the Gemini.",
+    allow_flagging="never"
+)
+
+if __name__ == "__main__":
+    demo.launch(debug=True)