{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Demo with Lightning⚡Flash: Image 🌹 Classification on TPU\n\n**REF: https://lightning-flash.readthedocs.io/en/stable/reference/image_classification.html**\n\nThe task of identifying what is in an image is called image classification. Typically, Image Classification is used to identify images containing a single object. The task predicts which ‘class’ the image most likely belongs to with a degree of certainty. A class is a label that describes what is in an image, such as ‘car’, ‘house’, ‘cat’ etc.","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"markdown","source":"## 0. Installing dependencies\n\nincluding this notebook you can pull packages for your offline kernels...\n\nSee: [Easy Kaggle Offline Submission With Chaining Kernel Notebooks](https://towardsdatascience.com/easy-kaggle-offline-submission-with-chaining-kernels-30bba5ea5c4d)","metadata":{}},{"cell_type":"code","source":"# !pip download -q \"icevision[all]\" 'lightning-flash[image]' --dest frozen_packages --prefer-binary\n# !pip download -q 'torchmetrics==0.7.*' --dest frozen_packages --prefer-binary\n# !pip download -q effdet timm segmentation-models-pytorch --dest frozen_packages --prefer-binary\n# !wget https://storage.googleapis.com/tpu-pytor -q -P frozen_packages/\n# !rm frozen_packages/torch-*\n# !ls -l frozen_packages | grep -e torch -e lightning -e timm","metadata":{"execution":{"iopub.status.busy":"2022-05-12T11:19:47.472692Z","iopub.execute_input":"2022-05-12T11:19:47.473083Z","iopub.status.idle":"2022-05-12T11:19:47.490638Z","shell.execute_reply.started":"2022-05-12T11:19:47.472954Z","shell.execute_reply":"2022-05-12T11:19:47.489629Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version 1.8 --apt-packages libomp5 libopenblas-dev\n!pip install -q -U 'lightning-flash[image]' --find-links frozen_packages # --no-index\n!pip uninstall -y -q torchtext","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-05-12T11:19:47.510314Z","iopub.execute_input":"2022-05-12T11:19:47.511116Z","iopub.status.idle":"2022-05-12T11:21:09.409185Z","shell.execute_reply.started":"2022-05-12T11:19:47.511076Z","shell.execute_reply":"2022-05-12T11:21:09.408448Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\n\nimport flash\nfrom flash.image import ImageClassificationData, ImageClassifier","metadata":{"execution":{"iopub.status.busy":"2022-05-12T11:21:09.410675Z","iopub.execute_input":"2022-05-12T11:21:09.410927Z","iopub.status.idle":"2022-05-12T11:21:29.428220Z","shell.execute_reply.started":"2022-05-12T11:21:09.410896Z","shell.execute_reply":"2022-05-12T11:21:29.427213Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os, glob\nimport tensorflow as tf \nfrom functools import partial\nfrom tqdm.auto import tqdm\nimport matplotlib.pyplot as plt\n\nAUTO = tf.data.experimental.AUTOTUNE # instructs the API to read from multiple files if available.\n\ndef decode_image(image_data, height: int = 512, width: int = 512):\n    image = tf.image.decode_jpeg(image_data, channels=3)\n    image = tf.cast(image, tf.float32) / 255.0\n    image = tf.reshape(image, [height, width, 3])\n    return image\n\n# for tfr in tfr_dataset:\n#     ex = tf.train.Example()\n#     ex.ParseFromString(tfr.numpy())\n#     spl = json.loads(MessageToJson(ex))['features']['feature']\n#     print(spl.keys())\n#     break\n\n    \ndef read_tfrecord(example, with_labels: bool = False):\n    tfrec_format = {\n        \"image\": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring\n        \"id\": tf.io.FixedLenFeature([], tf.string),  # shape [] means single element\n    }\n    if with_labels:\n        tfrec_format.update({\n            \"class\": tf.io.FixedLenFeature([], tf.int64),  # shape [] means single element\n        })\n    example = tf.io.parse_single_example(example, tfrec_format)\n    example['image'] = decode_image(example['image'])\n    if with_labels:\n        example['class'] = tf.cast(example['class'], tf.int32)\n    return example\n\n\ndef convert_dataset(path_in: str, path_out: str, sfolder: str = \"train\"):\n    fnames = glob.glob(os.path.join(path_in, sfolder, \"*.tfrec\"))\n    print(list(map(os.path.basename, fnames)))\n    # automatically interleaves reads from multiple files\n    dataset = tf.data.TFRecordDataset(fnames, num_parallel_reads=AUTO)\n    map_tfrecord = partial(read_tfrecord, with_labels=sfolder != \"test\")\n    dataset = dataset.map(map_tfrecord, num_parallel_calls=AUTO)\n\n    for spl in tqdm(dataset):\n        img_name = spl[\"id\"].numpy().decode(\"utf-8\")  + \".jpg\"\n        folders = [path_out, sfolder]\n        if \"class\" in spl:\n            folders.append(str(spl[\"class\"].numpy()))\n        img_path = os.path.join(*folders, img_name)\n        os.makedirs(os.path.dirname(img_path), exist_ok=True)\n        # print(img_path)\n        plt.imsave(img_path, spl[\"image\"].numpy())","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-05-12T11:21:29.429748Z","iopub.execute_input":"2022-05-12T11:21:29.430270Z","iopub.status.idle":"2022-05-12T11:21:29.445775Z","shell.execute_reply.started":"2022-05-12T11:21:29.430234Z","shell.execute_reply":"2022-05-12T11:21:29.445085Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"! rm -rf /kaggle/working/jpeg-512x512\n\nPATH_TFRECORD = \"/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512\"\nPATH_DATASET = \"/kaggle/working/jpeg-512x512\"\n\nconvert_dataset(PATH_TFRECORD, PATH_DATASET, \"train\")\nconvert_dataset(PATH_TFRECORD, PATH_DATASET, \"val\")\nconvert_dataset(PATH_TFRECORD, PATH_DATASET, \"test\")","metadata":{"execution":{"iopub.status.busy":"2022-05-12T11:21:29.447993Z","iopub.execute_input":"2022-05-12T11:21:29.448955Z","iopub.status.idle":"2022-05-12T11:30:52.748069Z","shell.execute_reply.started":"2022-05-12T11:21:29.448901Z","shell.execute_reply":"2022-05-12T11:30:52.746689Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 1. Create the DataModule","metadata":{}},{"cell_type":"code","source":"datamodule = ImageClassificationData.from_folders(\n    train_folder=os.path.join(PATH_DATASET, \"train\"),\n    val_folder=os.path.join(PATH_DATASET, \"val\"),\n    predict_folder=os.path.join(PATH_DATASET, \"test\"),\n    batch_size=12,\n    transform_kwargs={\"image_size\": (380, 380), \"mean\": (0.485, 0.456, 0.406), \"std\": (0.229, 0.224, 0.225)},\n    num_workers=3,\n)","metadata":{"execution":{"iopub.status.busy":"2022-05-12T11:30:52.750954Z","iopub.execute_input":"2022-05-12T11:30:52.751276Z","iopub.status.idle":"2022-05-12T11:30:53.011654Z","shell.execute_reply.started":"2022-05-12T11:30:52.751236Z","shell.execute_reply":"2022-05-12T11:30:53.010921Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\n# datamodule.show_train_batch()\n\nnb_samples = 9\nfig, axarr = plt.subplots(ncols=3, nrows=3, figsize=(8, 8))\n\nfor batch in datamodule.train_dataloader():\n    print(batch.keys())\n    for i, (img, lb) in enumerate(list(zip(batch[\"input\"], batch[\"target\"]))[:nb_samples]):\n        img = np.rollaxis(img.numpy(), 0, 3)\n        axarr[i % 3, i // 3].imshow(img, vmin=-5., vmax=5.)\n        axarr[i % 3, i // 3].set_title(lb)\n    break","metadata":{"execution":{"iopub.status.busy":"2022-05-12T11:30:53.013102Z","iopub.execute_input":"2022-05-12T11:30:53.013419Z","iopub.status.idle":"2022-05-12T11:30:55.417423Z","shell.execute_reply.started":"2022-05-12T11:30:53.013388Z","shell.execute_reply":"2022-05-12T11:30:55.416634Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2. Build the task","metadata":{}},{"cell_type":"code","source":"model = ImageClassifier(\n    backbone=\"tf_efficientnet_b4_ns\",\n    pretrained=True,\n    labels=datamodule.labels,\n)","metadata":{"execution":{"iopub.status.busy":"2022-05-12T11:36:38.620476Z","iopub.execute_input":"2022-05-12T11:36:38.620808Z","iopub.status.idle":"2022-05-12T11:36:39.108826Z","shell.execute_reply.started":"2022-05-12T11:36:38.620773Z","shell.execute_reply":"2022-05-12T11:36:39.107937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3. Create the trainer and finetune the model","metadata":{}},{"cell_type":"code","source":"from pytorch_lightning.loggers import CSVLogger\n\ntrainer = flash.Trainer(\n    max_epochs=5,\n    logger=CSVLogger(save_dir='logs/'),\n    precision=16,\n    tpu_cores=1,\n)","metadata":{"execution":{"iopub.status.busy":"2022-05-12T11:37:37.456296Z","iopub.execute_input":"2022-05-12T11:37:37.456763Z","iopub.status.idle":"2022-05-12T11:37:37.464990Z","shell.execute_reply.started":"2022-05-12T11:37:37.456726Z","shell.execute_reply":"2022-05-12T11:37:37.463673Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.finetune(model, datamodule=datamodule, strategy=\"freeze\")\n\n## Save the model!\ntrainer.save_checkpoint(\"image_classification_model.pt\")","metadata":{"execution":{"iopub.status.busy":"2022-05-12T11:37:39.149264Z","iopub.execute_input":"2022-05-12T11:37:39.150205Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sn\n\nmetrics = pd.read_csv(f'{trainer.logger.log_dir}/metrics.csv')\ndel metrics[\"step\"]\nmetrics.set_index(\"epoch\", inplace=True)\n# display(metrics.dropna(axis=1, how=\"all\").head())\ng = sn.relplot(data=metrics, kind=\"line\")\nplt.grid()","metadata":{"execution":{"iopub.status.busy":"2022-05-12T11:35:28.086810Z","iopub.status.idle":"2022-05-12T11:35:28.087137Z","shell.execute_reply.started":"2022-05-12T11:35:28.086982Z","shell.execute_reply":"2022-05-12T11:35:28.086999Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4. Predict what's on a few images!","metadata":{}},{"cell_type":"code","source":"predictions = trainer.predict(model, datamodule=datamodule, output=\"labels\")\nprint(predictions)","metadata":{"execution":{"iopub.status.busy":"2022-05-12T11:35:28.088292Z","iopub.status.idle":"2022-05-12T11:35:28.088603Z","shell.execute_reply.started":"2022-05-12T11:35:28.088441Z","shell.execute_reply":"2022-05-12T11:35:28.088457Z"},"trusted":true},"execution_count":null,"outputs":[]}]}