| import os |
| import json |
| from tqdm import tqdm |
| import openai |
| import traceback |
| import threading |
| from concurrent.futures import ThreadPoolExecutor |
| import json_repair |
|
|
| lock = threading.Lock() |
|
|
| client = openai.Client( |
| api_key = os.environ.get("DEEPSEEK_API_KEY", "[DEEPSEEK_API_KEY]"), |
| base_url = "https://api.deepseek.com/v1", |
| ) |
|
|
| with open("output1.json", "r") as f: |
| data = json.load(f) |
|
|
| new_data = {} |
| if os.path.exists("output.json"): |
| with open("output.json", "r") as f: |
| new_data = json.load(f) |
|
|
| progress = tqdm(data.items(), total=len(data), leave=True, position=0) |
| with ThreadPoolExecutor(max_workers=8) as executor: |
| for key, value in data.items(): |
| if key in new_data and new_data[key]: |
| progress.update(1) |
| continue |
| def process_and_update(key, value): |
| try: |
| prompt = "Given the following description of objects in an image, format the information as a JSON array where each object has 'name', 'bbox', and 'confidence' fields. The 'bbox' should be a list of four integers representing [x_min, y_min, width, height]. Here is the description: " + value + "\n\nOutput Format:\nOutput in JSON format: [{\"name\": object_name1, \"bbox\": bounding_box1, \"confidence\": confidence1}, {\"name\": object_name2, \"bbox\": bounding_box2, \"confidence\": confidence2}].\nDO NOT OUTPUT any other content besides JSON. If the name or bbox of certain objects cannot be determined, skip them in the output." |
| response = client.chat.completions.create( |
| model="deepseek-chat", |
| messages=[ |
| { |
| "role": "user", |
| "content": prompt |
| } |
| ], |
| temperature=0.1, |
| ) |
| resp_json = response.choices[0].message.content |
| resp_dict = json_repair.loads(resp_json) if resp_json else "" |
| new_data[key] = resp_dict |
| except Exception as e: |
| print(f"Error processing key {key}: {e}") |
| print(traceback.format_exc()) |
| new_data[key] = "" |
| finally: |
| progress.update(1) |
| with lock: |
| with open("output.json", "w") as f: |
| json.dump(new_data, f, indent=4, ensure_ascii=False) |
| executor.submit(process_and_update, key, value) |
|
|
