import argparse

import torch
from fvcore.nn import FlopCountAnalysis
from ultralytics import YOLO


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('model', type=str, help='Model path for validation.')
    parser.add_argument('--imgsz', default=928, type=int, help='Image size to validate.')
    args = parser.parse_args()

    model = torch.load(args.model, map_location='cpu')['model'].float()
    fca = FlopCountAnalysis(
        model=model.eval(),
        inputs=torch.rand(1, 3, args.imgsz, args.imgsz),
    )
    fca.unsupported_ops_warnings(False)
    fca.uncalled_modules_warnings(False)

    print(f"{fca.total() * 1e-9:.2f} GMACS ")