roqaia123 commited on
Commit
e421df7
·
verified ·
1 Parent(s): c898b20

Delete test.py

Browse files
Files changed (1) hide show
  1. test.py +0 -117
test.py DELETED
@@ -1,117 +0,0 @@
1
-
2
- import gradio as gr
3
- import os
4
- from io import BytesIO
5
- from PIL import Image, ImageDraw, ImageFont
6
- from PIL import ImageColor
7
- import json
8
- import google.generativeai as genai
9
- from google.generativeai import types
10
- from dotenv import load_dotenv
11
- from IPython.display import display
12
-
13
- # 1. SETUP API KEY
14
- # ----------------
15
- load_dotenv()
16
- api_key = os.getenv("Gemini_API_Key")
17
- # Configure the Google AI library
18
- genai.configure(api_key=api_key)
19
-
20
-
21
- # 2. DEFINE MODEL AND INSTRUCTIONS
22
-
23
- bounding_box_system_instructions = """
24
- Return bounding boxes as a JSON array with labels. Never return masks or code fencing. Limit to 25 objects.
25
- If an object is present multiple times, name them according to their unique characteristic (colors, size, position, unique characteristics, etc..).
26
- """
27
- model = genai.GenerativeModel( model_name='gemini-2.5-flash', system_instruction=bounding_box_system_instructions)
28
- generation_config = genai.types.GenerationConfig(
29
- temperature=0.5,
30
-
31
- )
32
-
33
-
34
-
35
- # 3. PREPARE IMAGE AND PROMPT
36
-
37
-
38
-
39
-
40
- def parse_json(json_output):
41
- lines = json_output.splitlines()
42
- for i, line in enumerate(lines):
43
- if line == "```json":
44
- json_output = "\n".join(lines[i+1:]) # Remove everything before "```json"
45
- json_output = json_output.split("```")[0] # Remove everything after the closing "```"
46
- break
47
- return json_output
48
- print("After parsing JSON from model response...")
49
-
50
-
51
- def plot_bounding_boxes(im, bounding_boxes):
52
- """
53
- Plots bounding boxes on an image with labels.
54
- """
55
- additional_colors = [colorname for (colorname, colorcode) in ImageColor.colormap.items()]
56
-
57
- im = im.copy()
58
- width, height = im.size
59
- draw = ImageDraw.Draw(im)
60
- colors = [
61
- 'red', 'green', 'blue', 'yellow', 'orange', 'pink', 'purple', 'cyan',
62
- 'lime', 'magenta', 'violet', 'gold', 'silver'
63
- ] + additional_colors
64
-
65
- try:
66
- # Use a default font if NotoSansCJK is not available
67
- try:
68
- font = ImageFont.load_default()
69
- except OSError:
70
- print("NotoSansCJK-Regular.ttc not found. Using default font.")
71
- font = ImageFont.load_default()
72
-
73
- bounding_boxes_json = json.loads(bounding_boxes)
74
- for i, bounding_box in enumerate(bounding_boxes_json):
75
- color = colors[i % len(colors)]
76
- abs_y1 = int(bounding_box["box_2d"][0] / 1000 * height)
77
- abs_x1 = int(bounding_box["box_2d"][1] / 1000 * width)
78
- abs_y2 = int(bounding_box["box_2d"][2] / 1000 * height)
79
- abs_x2 = int(bounding_box["box_2d"][3] / 1000 * width)
80
-
81
- if abs_x1 > abs_x2:
82
- abs_x1, abs_x2 = abs_x2, abs_x1
83
-
84
- if abs_y1 > abs_y2:
85
- abs_y1, abs_y2 = abs_y2, abs_y1
86
-
87
- # Draw bounding box and label
88
- draw.rectangle(((abs_x1, abs_y1), (abs_x2, abs_y2)), outline=color, width=4)
89
- if "label" in bounding_box:
90
- draw.text((abs_x1 + 8, abs_y1 + 6), bounding_box["label"], fill=color, font=font)
91
- except Exception as e:
92
- print(f"Error drawing bounding boxes: {e}")
93
-
94
- return im
95
-
96
- prompt = "Identify and label the objects in the image. Return only the JSON array of bounding boxes and labels as per the system instructions."
97
- image = "Images/cookies.jpg"
98
- img = Image.open(BytesIO(open(image, "rb").read()))
99
- print(f"Original image size: {img.size}")
100
-
101
- # resize the image to a max width of 1024 while maintaining aspect ratio
102
- im = Image.open(image).resize((1024, int(1024 * img.size[1] / img.size[0])), Image.Resampling.LANCZOS)
103
- print(f"Resized image size: {im.size}")
104
- im.show()
105
-
106
- # Run model to find bounding boxes
107
- response = model.generate_content([prompt, im], generation_config=generation_config)
108
- print("Raw model response:")
109
- print(response.text )
110
-
111
- bounding_boxes=parse_json(response.text)
112
-
113
-
114
- im_with_boxes = plot_bounding_boxes(im, bounding_boxes)
115
- display(im_with_boxes)
116
- im_with_boxes.save("output_imags/cookies_bounding_boxes.jpg")
117
- print("Bounding boxes plotted on image.")