tf_model_stats

import os
import tensorflow as tf
from spyglass.model import model_loader
from sales_zone.tf_models import ObjectDetector


def get_model_flops(model_cls, model_config: dict):
    with model_loader(model_cls,
                      model_config=model_config) as model:
        graph = model._session.graph

        flops = tf.profiler.profile(graph, run_meta=tf.RunMetadata(),
                                    cmd='op', options=tf.profiler.ProfileOptionBuilder.float_operation())
        return flops.total_float_ops


def format_model_flops(model_config: dict, flops: int) -> str:
    model_name = os.path.basename(model_config['weights'])
    return f"<{model_name}> : {flops} flops"


model_config = {"weights": "data/models/person_detector_from_above/resnet-fpn-768x512_top-view-obi_fp32.pb"}
print(format_model_flops(model_config, flops=get_model_flops(ObjectDetector, model_config)))


model_config = {"weights": "data/models/person_detector_from_side/mobilenet-fpn-768x512_ms-coco_tf32.pb"}
print(format_model_flops(model_config, flops=get_model_flops(ObjectDetector, model_config)))