import argparse
import numpy as np
from ultralytics import YOLO

LABELS = ['angry', 'disgust', 'fear', 'happy', 'neutral', 'sad', 'surprise']

def predict_probs(model, img_paths, imgsz, device=0):
    results = model.predict(img_paths, imgsz=imgsz, device=device, verbose=False)
    probs = []
    for r in results:
        probs.append(r.probs.data.cpu().numpy())  # shape (7,)
    return np.stack(probs, axis=0)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--image", required=True, help="Path to an image file")
    ap.add_argument("--device", default="0", help="CUDA device index or 'cpu'")
    args = ap.parse_args()

    mA = YOLO("weights/modelA_yolo11s_96.pt")
    mB = YOLO("weights/modelB_yolo11s_128.pt")

    pA = predict_probs(mA, [args.image], imgsz=96, device=args.device)[0]
    pB = predict_probs(mB, [args.image], imgsz=128, device=args.device)[0]

    p = (pA + pB) / 2.0
    pred = int(np.argmax(p))
    print("pred_label:", LABELS[pred])
    print("probs:", {LABELS[i]: float(p[i]) for i in range(len(LABELS))})

if __name__ == "__main__":
    main()