Jimmyzheng-10 commited on
Commit
871a905
·
1 Parent(s): 4796f79

Hide API key

Browse files
screencoder/.gitignore CHANGED
@@ -8,4 +8,4 @@ dist/
8
  .venv
9
 
10
  # API keys
11
- # doubao_api.txt
 
8
  .venv
9
 
10
  # API keys
11
+ doubao_api.txt
screencoder/block_parsor.py CHANGED
@@ -130,13 +130,19 @@ def resolve_containment(bboxes: dict[str, tuple[int, int, int, int]]) -> dict[st
130
  return {name: bbox for name, bbox in bboxes.items() if name not in removed}
131
 
132
  # sequential version of bbox parsing: Using recursive detection with mask
133
- def sequential_component_detection(image_path: str, api_path: str, temp_dir: str) -> dict[str, tuple[int, int, int, int]]:
134
  """
135
  Sequential processing flow: detect each component in turn, mask the image after each detection
136
  """
137
  bboxes = {}
138
  current_image_path = image_path
139
- ark_client = Doubao(api_path)
 
 
 
 
 
 
140
 
141
  image = cv2.imread(image_path)
142
  if image is None:
@@ -222,7 +228,7 @@ def main():
222
  tmp_dir = os.path.join(base_dir, 'data', 'tmp', run_id)
223
 
224
  image_path = os.path.join(tmp_dir, f"{run_id}.png")
225
- api_path = os.path.join(base_dir, "doubao_api.txt")
226
  json_output_path = os.path.join(tmp_dir, f"{run_id}_bboxes.json")
227
  annotated_image_output_path = os.path.join(tmp_dir, f"{run_id}_with_bboxes.png")
228
 
@@ -230,9 +236,9 @@ def main():
230
  print(f"Debug - base_dir: {base_dir}")
231
  print(f"Debug - tmp_dir: {tmp_dir}")
232
  print(f"Debug - image_path: {image_path}")
233
- print(f"Debug - api_path: {api_path}")
234
  print(f"Debug - image_path exists: {os.path.exists(image_path)}")
235
- print(f"Debug - api_path exists: {os.path.exists(api_path)}")
236
 
237
  # List contents of tmp_dir if it exists
238
  if os.path.exists(tmp_dir):
@@ -246,15 +252,17 @@ def main():
246
  save_bboxes_to_json({}, json_output_path)
247
  exit(1)
248
 
249
- if not os.path.exists(api_path):
250
- print(f"Error: API key file not found at {api_path}")
 
251
  # Create empty json file so the pipeline doesn't break
252
  save_bboxes_to_json({}, json_output_path)
253
  exit(1)
254
 
255
  print(f"--- Starting BBox Parsing for run_id: {run_id} ---")
256
 
257
- client = Doubao(api_path)
 
258
  base64_image = encode_image(image_path)
259
  if not base64_image:
260
  print(f"Error: Failed to encode image {image_path}")
 
130
  return {name: bbox for name, bbox in bboxes.items() if name not in removed}
131
 
132
  # sequential version of bbox parsing: Using recursive detection with mask
133
+ def sequential_component_detection(image_path: str, temp_dir: str) -> dict[str, tuple[int, int, int, int]]:
134
  """
135
  Sequential processing flow: detect each component in turn, mask the image after each detection
136
  """
137
  bboxes = {}
138
  current_image_path = image_path
139
+
140
+ # Check for API key - first try environment variable, then use provided path
141
+ api_key = os.environ.get('API_key')
142
+ if not api_key:
143
+ print(f"Error: API key not found in environment variable 'API_key'")
144
+ exit(1)
145
+ ark_client = Doubao(api_key)
146
 
147
  image = cv2.imread(image_path)
148
  if image is None:
 
228
  tmp_dir = os.path.join(base_dir, 'data', 'tmp', run_id)
229
 
230
  image_path = os.path.join(tmp_dir, f"{run_id}.png")
231
+ # api_path = os.path.join(base_dir, "doubao_api.txt")
232
  json_output_path = os.path.join(tmp_dir, f"{run_id}_bboxes.json")
233
  annotated_image_output_path = os.path.join(tmp_dir, f"{run_id}_with_bboxes.png")
234
 
 
236
  print(f"Debug - base_dir: {base_dir}")
237
  print(f"Debug - tmp_dir: {tmp_dir}")
238
  print(f"Debug - image_path: {image_path}")
239
+ # print(f"Debug - api_path: {api_path}")
240
  print(f"Debug - image_path exists: {os.path.exists(image_path)}")
241
+ # print(f"Debug - api_path exists: {os.path.exists(api_path)}")
242
 
243
  # List contents of tmp_dir if it exists
244
  if os.path.exists(tmp_dir):
 
252
  save_bboxes_to_json({}, json_output_path)
253
  exit(1)
254
 
255
+ api_key = os.environ.get('API_key')
256
+ if not api_key:
257
+ print(f"Error: API key not found in environment variable 'API_key'")
258
  # Create empty json file so the pipeline doesn't break
259
  save_bboxes_to_json({}, json_output_path)
260
  exit(1)
261
 
262
  print(f"--- Starting BBox Parsing for run_id: {run_id} ---")
263
 
264
+ # Use environment variable if available, otherwise use file path
265
+ client = Doubao(api_key)
266
  base64_image = encode_image(image_path)
267
  if not base64_image:
268
  print(f"Error: Failed to encode image {image_path}")
screencoder/html_generator.py CHANGED
@@ -299,12 +299,15 @@ def main():
299
 
300
  generate_html(root, output_html_path)
301
 
302
- api_path = os.path.join(base_dir, "doubao_api.txt")
303
- if not os.path.exists(api_path):
304
- print(f"Error: API key not found at {api_path}")
 
 
305
  exit(1)
306
 
307
- bot = Doubao(api_path, model="doubao-1.5-thinking-vision-pro-250428")
 
308
  code_dict = generate_code_parallel(root, img_path, bot, user_instruction)
309
  code_substitution(output_html_path, code_dict)
310
 
 
299
 
300
  generate_html(root, output_html_path)
301
 
302
+ # Check for API key - first try environment variable, then file
303
+ api_key = os.environ.get('API_key')
304
+ # api_path = os.path.join(base_dir, "doubao_api.txt")
305
+ if not api_key:
306
+ print(f"Error: API key not found in environment variable 'API_key'")
307
  exit(1)
308
 
309
+ # Use environment variable if available, otherwise use file path
310
+ bot = Doubao(api_key, model="doubao-1.5-thinking-vision-pro-250428")
311
  code_dict = generate_code_parallel(root, img_path, bot, user_instruction)
312
  code_substitution(output_html_path, code_dict)
313