blakeblackshear.frigate/notebooks/YOLO_NAS_Pretrained_Export.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rmuF9iKWTbdk"
      },
      "outputs": [],
      "source": [
        "! pip install -q git+https://github.com/Deci-AI/super-gradients.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NiRCt917KKcL"
      },
      "outputs": [],
      "source": [
        "! sed -i 's/sghub.deci.ai/sg-hub-nv.s3.amazonaws.com/' /usr/local/lib/python3.11/dist-packages/super_gradients/training/pretrained_models.py\n",
        "! sed -i 's/sghub.deci.ai/sg-hub-nv.s3.amazonaws.com/' /usr/local/lib/python3.11/dist-packages/super_gradients/training/utils/checkpoint_utils.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dTB0jy_NNSFz"
      },
      "outputs": [],
      "source": [
        "from super_gradients.common.object_names import Models\n",
        "from super_gradients.conversion import DetectionOutputFormatMode\n",
        "from super_gradients.training import models\n",
        "\n",
        "model = models.get(Models.YOLO_NAS_S, pretrained_weights=\"coco\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GymUghyCNXem"
      },
      "outputs": [],
      "source": [
        "# export the model for compatibility with Frigate\n",
        "\n",
        "model.export(\"yolo_nas_s.onnx\",\n",
        "             output_predictions_format=DetectionOutputFormatMode.FLAT_FORMAT,\n",
        "             max_predictions_per_image=20,\n",
        "             num_pre_nms_predictions=300,\n",
        "             confidence_threshold=0.4,\n",
        "             input_image_shape=(320,320),\n",
        "            )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uBhXV5g4Nh42"
      },
      "outputs": [],
      "source": [
        "from google.colab import files\n",
        "\n",
        "files.download('yolo_nas_s.onnx')"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}