File size: 2,466 Bytes
1da285f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | 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)
|