File size: 1,095 Bytes
159b2dd | 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 | 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()
|