Labels : retry if name isn't known

master
Sébastien Miquel 2026-02-26 08:48:25 +01:00
parent 63489c95c7
commit f384121e9b
1 changed files with 10 additions and 0 deletions

View File

@ -205,6 +205,9 @@ for path_str in args.input_paths:
labels_txt = (INPUT_DIR / "labels").read_text() labels_txt = (INPUT_DIR / "labels").read_text()
valid_labels_set = set(line.strip() for line in labels_txt.splitlines() if line.strip()) valid_labels_set = set(line.strip() for line in labels_txt.splitlines() if line.strip())
names_txt = (INPUT_DIR / "names").read_text() names_txt = (INPUT_DIR / "names").read_text()
valid_names_set = set(line.strip() for line in names_txt.splitlines() if line.strip())
valid_names_set.add("Unknown")
client = genai.Client(api_key=api_key) client = genai.Client(api_key=api_key)
# Group files by Copy ID (e.g. Copie01_01.jpg -> Copie01) # Group files by Copy ID (e.g. Copie01_01.jpg -> Copie01)
@ -265,12 +268,19 @@ def process_copy_group(group_key, files):
annota = AnnotationData.model_validate_json(response.text) annota = AnnotationData.model_validate_json(response.text)
unknown = [item.label for item in annota.list if item.label not in valid_labels_set] unknown = [item.label for item in annota.list if item.label not in valid_labels_set]
name = annota.name
if unknown: if unknown:
print(f"Error: {image_file.name} contained unknown labels: {unknown}") print(f"Error: {image_file.name} contained unknown labels: {unknown}")
if attempt == 0: if attempt == 0:
print("Retrying request...") print("Retrying request...")
continue # Retry immediately continue # Retry immediately
if name not in valid_names_set:
print(f"Error: {image_file.name} returned unknown name : {name}")
if attempt == 0:
print("Retrying request...")
continue # Retry immediately
# Save result # Save result
with open(output_json, "w", encoding="utf-8") as f: with open(output_json, "w", encoding="utf-8") as f:
json.dump(annota.model_dump(), f, indent=2) json.dump(annota.model_dump(), f, indent=2)