# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] id="XjQxmm04iTx4" tags=[]
# <a href="https://colab.research.google.com/github/open-mmlab/mmclassification/blob/master/docs/tutorials/MMClassification_tools.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + [markdown] id="4z0JDgisPRr-"
# # MMClassification tools tutorial on Colab
#
# In this tutorial, we will introduce the following content:
#
# * How to install MMCls
# * Prepare data
# * Prepare the config file
# * Train and test model with shell command

# + [markdown] id="inm7Ciy5PXrU"
# ## Install MMClassification
#
# Before using MMClassification, we need to prepare the environment with the following steps:
#
# 1. Install Python, CUDA, C/C++ compiler and git
# 2. Install PyTorch (CUDA version)
# 3. Install mmcv
# 4. Clone mmcls source code from GitHub and install it
#
# Because this tutorial is on Google Colab, and the basic environment has been completed, we can skip the first two steps.

# + [markdown] id="TDOxbcDvPbNk"
# ### Check environment

# + colab={"base_uri": "https://localhost:8080/"} id="c6MbAw10iUJI" outputId="8d3d6b53-c69b-4425-ce0c-bfb8d31ab971"
# %cd /content

# + colab={"base_uri": "https://localhost:8080/"} id="4IyFL3MaiYRu" outputId="c46dc718-27de-418b-da17-9d5a717e8424"
# !pwd

# + colab={"base_uri": "https://localhost:8080/"} id="DMw7QwvpiiUO" outputId="0d852285-07c4-48d3-e537-4a51dea04d10"
# Check nvcc version
# !nvcc -V

# + colab={"base_uri": "https://localhost:8080/"} id="4VIBU7Fain4D" outputId="fb34a7b6-8eda-4180-e706-1bf67d1a6fd4"
# Check GCC version
# !gcc --version

# + colab={"base_uri": "https://localhost:8080/"} id="24lDLCqFisZ9" outputId="304ad2f7-a9bb-4441-d25b-09b5516ccd74"
# Check PyTorch installation
import torch, torchvision
print(torch.__version__)
print(torch.cuda.is_available())

# + [markdown] id="R2aZNLUwizBs"
# ### Install MMCV
#
# MMCV is the basic package of all OpenMMLab packages. We have pre-built wheels on Linux, so we can download and install them directly.
#
# Please pay attention to PyTorch and CUDA versions to match the wheel.
#
# In the above steps, we have checked the version of PyTorch and CUDA, and they are 1.9.0 and 11.1 respectively, so we need to choose the corresponding wheel.
#
# In addition, we can also install the full version of mmcv (mmcv-full). It includes full features and various CUDA ops out of the box, but needs a longer time to build.

# + colab={"base_uri": "https://localhost:8080/"} id="nla40LrLi7oo" outputId="a17d50d6-05b7-45d6-c3fb-6a2507415cf5"
# Install mmcv
# !pip install mmcv -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.9.0/index.html
# # !pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.9.0/index.html

# + [markdown] id="GDTUrYvXjlRb"
# ### Clone and install MMClassification
#
# Next, we clone the latest mmcls repository from GitHub and install it.

# + colab={"base_uri": "https://localhost:8080/"} id="Bwme6tWHjl5s" outputId="7e2d54c8-b134-405a-b014-194da1708776"
# Clone mmcls repository
# !git clone https://github.com/open-mmlab/mmclassification.git
# %cd mmclassification/

# Install MMClassification from source
# !pip install -e . 

# + colab={"base_uri": "https://localhost:8080/"} id="hFg_oSG4j3zB" outputId="1cc74bac-f918-4f0e-bf56-9f13447dfce1"
# Check MMClassification installation
import mmcls
print(mmcls.__version__)

# + [markdown] id="HCOHRp3iV5Xk"
# ## Prepare data

# + colab={"base_uri": "https://localhost:8080/"} id="XHCHnKb_Qd3P" outputId="35496010-ee57-4e72-af00-2af55dc80f47"
# Download the dataset (cats & dogs dataset)
# !wget https://www.dropbox.com/s/wml49yrtdo53mie/cats_dogs_dataset_reorg.zip?dl=0 -O cats_dogs_dataset.zip
# !mkdir -p data
# !unzip -q cats_dogs_dataset.zip -d ./data/

# + [markdown] id="e4t2P2aTQokX"
# **After downloading and extraction,** we get "Cats and Dogs Dataset" and the file structure is as below:
# ```
# data/cats_dogs_dataset
# ├── classes.txt
# ├── test.txt
# ├── val.txt
# ├── training_set
# │   ├── training_set
# │   │   ├── cats
# │   │   │   ├── cat.1.jpg
# │   │   │   ├── cat.2.jpg
# │   │   │   ├── ...
# │   │   ├── dogs
# │   │   │   ├── dog.2.jpg
# │   │   │   ├── dog.3.jpg
# │   │   │   ├── ...
# ├── val_set
# │   ├── val_set
# │   │   ├── cats
# │   │   │   ├── cat.3.jpg
# │   │   │   ├── cat.5.jpg
# │   │   │   ├── ...
# │   │   ├── dogs
# │   │   │   ├── dog.1.jpg
# │   │   │   ├── dog.6.jpg
# │   │   │   ├── ...
# ├── test_set
# │   ├── test_set
# │   │   ├── cats
# │   │   │   ├── cat.4001.jpg
# │   │   │   ├── cat.4002.jpg
# │   │   │   ├── ...
# │   │   ├── dogs
# │   │   │   ├── dog.4001.jpg
# │   │   │   ├── dog.4002.jpg
# │   │   │   ├── ...
# ```
#
# You can use shell command `tree data/cats_dogs_dataset` to check the structure.

# + colab={"base_uri": "https://localhost:8080/", "height": 297} id="46tyHTdtQy_Z" outputId="6124a89e-03eb-4917-a0bf-df6a391eb280"
# Pick an image and visualize it
from PIL import Image
Image.open('data/cats_dogs_dataset/training_set/training_set/cats/cat.1.jpg')

# + [markdown] id="My5Z6p7pQ3UC"
# ### Support new dataset
#
# We have two methods to support a new dataset in MMClassification.
#
# The simplest method is to re-organize the new dataset as the format of a dataset supported officially (like ImageNet). And the other method is to create a new dataset class, and more details are in [the docs](https://mmclassification.readthedocs.io/en/latest/tutorials/new_dataset.html#an-example-of-customized-dataset).
#
# In this tutorial, for convenience, we have re-organized the cats & dogs dataset as the format of ImageNet.

# + [markdown] id="P335gKt9Q5U-"
# Besides image files, it also includes the following files:
#
# 1. A class list file, and every line is a class.
#     ```
#     cats
#     dogs
#     ```
# 2. Training / Validation / Test annotation files. And every line includes an file path and the corresponding label.
#
#     ```
#     ...
#     cats/cat.3769.jpg 0
#     cats/cat.882.jpg 0
#     ...
#     dogs/dog.3881.jpg 1
#     dogs/dog.3377.jpg 1
#     ...
#     ```

# + [markdown] id="BafQ7ijBQ8N_"
# ## Train and test model with shell commands
#
# You can use shell commands provided by MMClassification to do the following task:
#
# 1. Train a model
# 2. Fine-tune a model
# 3. Test a model
# 4. Inference with a model
#
# The procedure to train and fine-tune a model is almost the same. And we have introduced how to do these tasks with Python API. In the following, we will introduce how to do them with shell commands. More details are in [the docs](https://mmclassification.readthedocs.io/en/latest/getting_started.html).

# + [markdown] id="Aj5cGMihURrZ"
# ### Fine-tune a model
#
# The steps to fine-tune a model are as below:
#
# 1. Prepare the custom dataset.
# 2. Create a new config file of the task.
# 3. Start training task by shell commands.
#
# We have finished the first step, and then we will introduce the next two steps.
#

# + [markdown] id="WBBV3aG79ZH5"
# #### Create a new config file
#
# To reuse the common parts of different config files, we support inheriting multiple base config files. For example, to fine-tune a MobileNetV2 model, the new config file can create the model's basic structure by inheriting `configs/_base_/models/mobilenet_v2_1x.py`.
#
# According to the common practice, we usually split whole configs into four parts: model, dataset, learning rate schedule, and runtime. Configs of each part are saved into one file in the `configs/_base_` folder. 
#
# And then, when creating a new config file, we can select some parts to inherit and only override some different configs.
#
# The head of the final config file should look like:
#
# ```python
# _base_ = [
#     '../_base_/models/mobilenet_v2_1x.py',
#     '../_base_/schedules/imagenet_bs256_epochstep.py',
#     '../_base_/default_runtime.py'
# ]
# ```
#
# Here, because the dataset configs are almost brand new, we don't need to inherit any dataset config file.
#
# Of course, you can also create an entire config file without inheritance, like `configs/mnist/lenet5.py`.

# + [markdown] id="_UV3oBhLRG8B"
# After that, we only need to set the part of configs we want to modify, because the inherited configs will be merged to the final configs.

# + colab={"base_uri": "https://localhost:8080/"} id="8QfM4qBeWIQh" outputId="a826f0cf-2633-4a9a-e49b-4be7eca5e3a0"
# %%writefile configs/mobilenet_v2/mobilenet_v2_1x_cats_dogs.py
_base_ = [
    '../_base_/models/mobilenet_v2_1x.py',
    '../_base_/schedules/imagenet_bs256_epochstep.py',
    '../_base_/default_runtime.py'
]

# ---- Model configs ----
# Here we use init_cfg to load pre-trained model.
# In this way, only the weights of backbone will be loaded.
# And modify the num_classes to match our dataset.

model = dict(
    backbone=dict(
        init_cfg = dict(
            type='Pretrained', 
            checkpoint='https://download.openmmlab.com/mmclassification/v0/mobilenet_v2/mobilenet_v2_batch256_imagenet_20200708-3b2dc3af.pth', 
            prefix='backbone')
    ),
    head=dict(
        num_classes=2,
        topk = (1, )
    ))

# ---- Dataset configs ----
# We re-organized the dataset as ImageNet format.
dataset_type = 'ImageNet'
img_norm_cfg = dict(
     mean=[124.508, 116.050, 106.438],
     std=[58.577, 57.310, 57.437],
     to_rgb=True)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='RandomResizedCrop', size=224, backend='pillow'),
    dict(type='RandomFlip', flip_prob=0.5, direction='horizontal'),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='ImageToTensor', keys=['img']),
    dict(type='ToTensor', keys=['gt_label']),
    dict(type='Collect', keys=['img', 'gt_label'])
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='Resize', size=(256, -1), backend='pillow'),
    dict(type='CenterCrop', crop_size=224),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='ImageToTensor', keys=['img']),
    dict(type='Collect', keys=['img'])
]
data = dict(
    # Specify the batch size and number of workers in each GPU.
    # Please configure it according to your hardware.
    samples_per_gpu=32,
    workers_per_gpu=2,
    # Specify the training dataset type and path
    train=dict(
        type=dataset_type,
        data_prefix='data/cats_dogs_dataset/training_set/training_set',
        classes='data/cats_dogs_dataset/classes.txt',
        pipeline=train_pipeline),
    # Specify the validation dataset type and path
    val=dict(
        type=dataset_type,
        data_prefix='data/cats_dogs_dataset/val_set/val_set',
        ann_file='data/cats_dogs_dataset/val.txt',
        classes='data/cats_dogs_dataset/classes.txt',
        pipeline=test_pipeline),
    # Specify the test dataset type and path
    test=dict(
        type=dataset_type,
        data_prefix='data/cats_dogs_dataset/test_set/test_set',
        ann_file='data/cats_dogs_dataset/test.txt',
        classes='data/cats_dogs_dataset/classes.txt',
        pipeline=test_pipeline))

# Specify evaluation metric
evaluation = dict(metric='accuracy', metric_options={'topk': (1, )})

# ---- Schedule configs ----
# Usually in fine-tuning, we need a smaller learning rate and less training epochs.
# Specify the learning rate
optimizer = dict(type='SGD', lr=0.005, momentum=0.9, weight_decay=0.0001)
optimizer_config = dict(grad_clip=None)
# Set the learning rate scheduler
lr_config = dict(policy='step', step=1, gamma=0.1)
runner = dict(type='EpochBasedRunner', max_epochs=2)

# ---- Runtime configs ----
# Output training log every 10 iterations.
log_config = dict(interval=10)

# + [markdown] id="chLX7bL3RP2F"
# #### Use shell command to start fine-tuning
#
# We use `tools/train.py` to fine-tune a model:
#
# ```shell
# python tools/train.py ${CONFIG_FILE} [optional arguments]
# ```
#
# And if you want to specify another folder to save log files and checkpoints, use the argument `--work_dir ${YOUR_WORK_DIR}`.
#
# If you want to ensure reproducibility, use the argument `--seed ${SEED}` to set a random seed. And the argument `--deterministic` can enable the deterministic option in cuDNN to further ensure reproducibility, but it may reduce the training speed.
#
# Here we use the `MobileNetV2` model and cats & dogs dataset as an example:
#

# + colab={"base_uri": "https://localhost:8080/"} id="gbFGR4SBRUYN" outputId="3412752c-433f-43c5-82a9-3495d1cd797a"
# !python tools/train.py \
#   configs/mobilenet_v2/mobilenet_v2_1x_cats_dogs.py \
#   --work-dir work_dirs/mobilenet_v2_1x_cats_dogs \
#   --seed 0 \
#   --deterministic

# + [markdown] id="m_ZSkwB5Rflb"
# ### Test a model
#
# We use `tools/test.py` to test a model:
#
# ```
# python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} [optional arguments]
# ```
#
# Here are some optional arguments:
#
# - `--metrics`: The evaluation metrics. The available choices are defined in the dataset class. Usually, you can specify "accuracy" to metric a single-label classification task.
# - `--metric-options`: The extra options passed to metrics. For example, by specifying "topk=1", the "accuracy" metric will calculate top-1 accuracy.
#
# More details are in the help docs of `tools/test.py`.
#
# Here we still use the `MobileNetV2` model we fine-tuned.

# + colab={"base_uri": "https://localhost:8080/"} id="Zd4EM00QRtyc" outputId="8788264f-83df-4419-9748-822c20538aa7"
# !python tools/test.py configs/mobilenet_v2/mobilenet_v2_1x_cats_dogs.py work_dirs/mobilenet_v2_1x_cats_dogs/latest.pth --metrics accuracy --metric-options topk=1

# + [markdown] id="IwThQkjaRwF7"
# ### Inference with a model
#
# Sometimes we want to save the inference results on a dataset, just use the command below.
#
# ```shell
# python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} [--out ${RESULT_FILE}]
# ```
#
# Arguments：
#
# - `--out`: The output filename. If not specified, the inference results won't be saved. It supports json, pkl and yml.
# - `--out-items`: What items will be saved. You can choose some of "class_scores", "pred_score", "pred_label" and "pred_class", or use "all" to select all of them.
#
# These items mean:
# - `class_scores`: The score of every class for each sample.
# - `pred_score`: The score of predict class for each sample.
# - `pred_label`: The label of predict class for each sample. It will read the label string of each class from the model, if the label strings are not saved, it will use ImageNet labels.
# - `pred_class`: The id of predict class for each sample. It's a group of integers. 
# - `all`: Save all items above.
# - `none`: Don't save any items above. Because the output file will save the metric besides inference results. If you want to save only metrics, you can use this option to reduce the output file size.
#

# + colab={"base_uri": "https://localhost:8080/"} id="6GVKloPHR0Fn" outputId="4f4cd414-1be6-4e17-985f-6449b8a3d9e8"
# !python tools/test.py configs/mobilenet_v2/mobilenet_v2_1x_cats_dogs.py work_dirs/mobilenet_v2_1x_cats_dogs/latest.pth --out results.json --out-items all

# + [markdown] id="G0NJI1s6e3FD"
# All inference results are saved in the output json file, and you can read it.

# + colab={"base_uri": "https://localhost:8080/", "height": 370} id="HJdJeLUafFhX" outputId="7614d546-7c2f-4bfd-ce63-4c2a5228620f"
import json

with open("./results.json", 'r') as f:
    results = json.load(f)

# Show the inference result of the first image.
print('class_scores:', results['class_scores'][0])
print('pred_class:', results['pred_class'][0])
print('pred_label:', results['pred_label'][0])
print('pred_score:', results['pred_score'][0])
Image.open('data/cats_dogs_dataset/training_set/training_set/cats/cat.1.jpg')

# + [markdown] id="1bEUwwzcVG8o"
# You can also use the visualization API provided by MMClassification to show the inference result.

# + id="BcSNyvAWRx20" colab={"base_uri": "https://localhost:8080/", "height": 304} outputId="0d68077f-2ec8-4f3d-8aaa-18d4021ca77b"
from mmcls.core.visualization import imshow_infos

filepath = 'data/cats_dogs_dataset/training_set/training_set/cats/cat.1.jpg'

result = {
    'pred_class': results['pred_class'][0],
    'pred_label': results['pred_label'][0],
    'pred_score': results['pred_score'][0],
}

img = imshow_infos(filepath, result)

StarcoderdataJupyterScriptsDedupFiltered