import json

from ultralytics import YOLO

MODELS: list[dict[str, str]] = [
    {"name": "yolov8n.pt"},
    {"name": "yolo26n.pt"},
    {"name": "yolo11n.pt"},
    {"name": "yolo10n.pt"},
    {"name": "yolo9n.pt"},
    {"name": "yolo3n.pt"},
]
PATH_SAVE_DATA: str = "othet.json"


def save_data(data: dict, model_name: str) -> None:
    with open(file=f"{model_name}.json", mode="w", encoding="utf-8") as file:
        json.dump(obj=data, fp=file, ensure_ascii=False, indent=2)


def analystic(model_data: dict[str, str]) -> None:
    model_name: str | None = model_data["name"]
    if model_name:
        model = YOLO(model=model_name)
        data: dict[str, str | list | dict] = {
            "model": model_name,
            "objects": [],
            "info": {},
        }
        result = model(source="input_image.jpg")
        for r in result:
            for box in r.boxes:
                class_id = int(box.cls[0])
                confidence = float(box.conf[0])
                class_name = model.names[class_id]
                if data["info"].get(class_name) is None:
                    data["info"][class_name] = {"count": 1, "confidence": confidence}
                else:
                    data["info"][class_name]["count"] += 1
                    data["info"][class_name]["confidence"] += confidence
                if isinstance(data["objects"], list):
                    data["objects"].append(
                        {
                            "class_id": class_id,
                            "class_name": class_name,
                            "confidence": confidence,
                        }
                    )

        result[0].save(f"output_{model_name}.jpg")
        save_data(data=data, model_name=model_name)


if __name__ == "__main__":
    for model in MODELS:
        analystic(model_data=model)