Get YOLO classes

get_yolo_classes.ipynb

{
  "cells": [
    {
      "metadata": {
        "trusted": true
      },
      "cell_type": "code",
      "source": "import torch",
      "execution_count": 1,
      "outputs": []
    },
    {
      "metadata": {
        "trusted": true
      },
      "cell_type": "code",
      "source": "model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5m, yolov5l, yolov5x, custom",
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": "Using cache found in /Users/spb/.cache/torch/hub/ultralytics_yolov5_master\nYOLOv5 🚀 2022-2-14 torch 1.8.0 CPU\n\nFusing layers... \nModel Summary: 213 layers, 7225885 parameters, 0 gradients\nAdding AutoShape... \n",
          "name": "stderr"
        }
      ]
    },
    {
      "metadata": {
        "trusted": true
      },
      "cell_type": "code",
      "source": "imgs = ['https://ultralytics.com/images/zidane.jpg',\n        'https://p0.piqsels.com/preview/1022/87/796/tree-plant-green-leaf-thumbnail.jpg',\n        'https://c4.wallpaperflare.com/wallpaper/996/421/193/boeing-747-400-klm-royal-dutch-airlines-hd-wallpaper-thumb.jpg']",
      "execution_count": 3,
      "outputs": []
    },
    {
      "metadata": {
        "trusted": true
      },
      "cell_type": "code",
      "source": "results = model(imgs)\nresults.print()",
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": "image 1/3: 720x1280 2 persons, 1 tie\nimage 2/3: 320x480 1 cat\nimage 3/3: 310x465 1 airplane\nSpeed: 157.0ms pre-process, 90.9ms inference, 2.6ms NMS per image at shape (3, 3, 448, 640)\n",
          "name": "stderr"
        }
      ]
    },
    {
      "metadata": {
        "trusted": true
      },
      "cell_type": "code",
      "source": "classes = [{results.names[int(p)] for p in pred[:, -1]} for pred in results.pred]\nclasses",
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "execution_count": 5,
          "data": {
            "text/plain": "[{'person', 'tie'}, {'cat'}, {'airplane'}]"
          },
          "metadata": {}
        }
      ]
    }
  ],
  "metadata": {
    "hide_input": false,
    "kernelspec": {
      "name": "conda-env-torch-py",
      "display_name": "Python [conda env:torch]",
      "language": "python"
    },
    "language_info": {
      "name": "python",
      "version": "3.7.11",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    },
    "gist": {
      "id": "",
      "data": {
        "description": "Get YOLO classes",
        "public": true
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}