ModOAP_Layout_Parser_Entrainement.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ModOAP_Layout_Parser_Entrainement.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pbGanx_tVzjv"
      },
      "source": [
        "# ModOAP - Entraînement à la détection de mise en page - Layout Parser\n",
        "\n",
        "Ce script permet d'entraîner un modèle de détection de mise en page de documents grâce à l'outil Layout Parser. \n",
        "\n",
        "L'entraînement doit être réalisé à partir d'un corpus de pages de documents annotées au format COCO.\n",
        "\n",
        "Le modèle entraîné peut être sauvegardé dans un dossier Google Drive, et utilisé pour la détection avec Layout Parser.\n",
        "\n",
        "Une fonctionnalité permet d'essayer un modèle entraîné sur une page de document et de visualiser les résultats.\n",
        "\n",
        "**Ce script doit être lancé dans un environnement d'exécution GPU : Exécution -> Modifier le type d'exécution -> GPU**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VfbWFXAUH0l4"
      },
      "source": [
        "#@title Imports et synchronisation avec Drive { display-mode: \"form\" }\n",
        "\n",
        "from google.colab import drive\n",
        "import os\n",
        "\n",
        "# chargement d'un google drive\n",
        "if not os.path.exists(\"/content/drive/MyDrive/\") :\n",
        "  drive.mount('/content/drive/')\n",
        "\n",
        "import shutil\n",
        "import requests\n",
        "import glob\n",
        "import json\n",
        "import cv2\n",
        "\n",
        "try:\n",
        "  import layoutparser as lp\n",
        "except ModuleNotFoundError:\n",
        "  !pip -q install 'git+https://github.com/facebookresearch/detectron2.git@v0.4#egg=detectron2' \n",
        "  !pip -q install -U layoutparser\n",
        "  import layoutparser as lp\n",
        "\n",
        "!git clone --quiet https://github.com/Layout-Parser/layout-model-training.git\n",
        "!wget -q -O ./layout-model-training/config_LayoutParser_PrimaDataset_original.yaml https://www.dropbox.com/s/yc92x97k50abynt/config.yaml?dl=1\n",
        "\n",
        "import torch\n",
        "torch.cuda.empty_cache()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TNcX1fhaVgR_"
      },
      "source": [
        "# Entrainement"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5g2wKSECCC0E",
        "cellView": "form"
      },
      "source": [
        "#@markdown ## Entrainement d'un modèle à partir d'un dataset COCO\n",
        "#@markdown Renseigner les paramètres avant de lancer la cellule\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown #### Entrer le chemin vers le dossier contenant le corpus d'entraînement au format COCO :\n",
        "chemin_corpus = \"\" #@param {type:\"string\"}\n",
        "#@markdown Exemple de chemin: /content/drive/MyDrive/Datasets/corpus_entrainement/\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown #### Entrer le chemin vers un dossier où sauvegarder le modèle entraîné  :\n",
        "chemin_destination_modele = \"\" #@param {type:\"string\"}\n",
        "#@markdown Exemple de chemin: /content/drive/MyDrive/Modeles/modele/\n",
        "\n",
        "#@markdown Le dossier est créé si non-existant\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown #### Spécifier les paramètres d'entraînement :\n",
        "\n",
        "#@markdown Nombre d'itérations :\n",
        "MAX_ITER =  300#@param {type:\"integer\"}\n",
        "#@markdown Fréquence de sauvegarde du modèle en nombre d'itérations :\n",
        "CHECKPOINT_PERIOD =  50#@param {type:\"integer\"}\n",
        "#@markdown Taux d'apprentissage (décimal compris entre 0 et 1) :\n",
        "LEARNING_RATE = 0.00025 #@param {type:\"number\"}\n",
        "#@markdown Apprentissage à partir d'un modèle préalablement entraîné :\n",
        "\n",
        "#@markdown Entrer le chemin vers le modèle .pth ou laisser vide pour partir du modèle pré-entraîné PRIMA\n",
        "MODELE = \"\" #@param {type:\"string\"}\n",
        "#@markdown Exemple de chemin: /content/drive/MyDrive/Modeles/modele/model_final.pth\n",
        "\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "\n",
        "if not os.path.exists(chemin_destination_modele) :\n",
        "  os.makedirs(chemin_destination_modele)\n",
        "\n",
        "# Transformation du fichier configuration pour l'entrainement du modele\n",
        "\n",
        "\n",
        "with open(\"/content/layout-model-training/config_LayoutParser_PrimaDataset_original.yaml\", \"r\") as c :\n",
        "  conf = c.read()\n",
        "conf = conf.replace(\"MAX_ITER: 60000\",\"MAX_ITER: {}\".format(MAX_ITER))\n",
        "conf = conf.replace(\"CHECKPOINT_PERIOD: 5000\",\"CHECKPOINT_PERIOD: {}\".format(CHECKPOINT_PERIOD)) \n",
        "conf = conf.replace(\"BASE_LR: 0.00025\",\"BASE_LR: {}\".format(LEARNING_RATE))\n",
        "if MODELE :\n",
        "  conf = conf.replace(\"WEIGHTS: https://www.dropbox.com/s/h7th27jfv19rxiy/model_final.pth?dl=1\",\"WEIGHTS: {}\".format(MODELE))\n",
        "with open(\"/content/layout-model-training/config_LayoutParser_PrimaDataset.yaml\", \"w\") as c :\n",
        "  c.write(conf)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# Lancement de l'entraînement\n",
        "\n",
        "ds_name = chemin_destination_modele.split(\"/\")[-1]\n",
        "annot_train = os.path.join(chemin_corpus,\"train\",\"*.json\")\n",
        "img_train = os.path.join(chemin_corpus,\"train\")\n",
        "annot_val = os.path.join(chemin_corpus,\"val\",\"*.json\")\n",
        "img_val = os.path.join(chemin_corpus,\"val\")\n",
        "\n",
        "\n",
        "!python /content/layout-model-training/tools/train_net.py --dataset_name $ds_name \\\n",
        "    --json_annotation_train $annot_train \\\n",
        "    --image_path_train      $img_train \\\n",
        "    --json_annotation_val   $annot_val \\\n",
        "    --image_path_val        $img_val \\\n",
        "    --config-file           /content/layout-model-training/config_LayoutParser_PrimaDataset.yaml \\\n",
        "    OUTPUT_DIR  $chemin_destination_modele \\\n",
        "    SOLVER.IMS_PER_BATCH 2\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0JxdxA7K6Ceh"
      },
      "source": [
        "# Essai d'un modèle sur un petit corpus"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Oa3ln2dUWyzb",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "cellView": "form",
        "outputId": "d6af23f5-e367-4258-e79e-57342f6ca64b"
      },
      "source": [
        "#@markdown ## Essai du modèle \n",
        "#@markdown Renseigner les paramètres avant de lancer la cellule\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown #### Entrer le chemin du modèle .pth à charger :\n",
        "modele = \"\" #@param {type:\"string\"}\n",
        "#@markdown Exemple de chemin: /content/drive/MyDrive/Modeles/modele/modele_final.pth\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown #### Entrer le chemin du fichier de configuration .yaml de ce modèle  :\n",
        "config = \"\" #@param {type:\"string\"}\n",
        "#@markdown Exemple de chemin: /content/drive/MyDrive/Modeles/modele/config.yaml\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown #### Entrer l'association id:nom pour chaque classe à détecter, séparées par un esapce (ou laisser vide et ne pas associer):\n",
        "\n",
        "labels = \"\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Exemple : 0:illustration 1:legende 2:titre\n",
        "\n",
        "try :\n",
        "  if labels : \n",
        "    dictlabels = {}\n",
        "    for asso in labels.split(\" \") :\n",
        "      dictlabels[asso.split(\":\")[0]] = asso.split(\":\")[1]\n",
        "\n",
        "    model = lp.Detectron2LayoutModel(config_path= config,\n",
        "                                  model_path=modele,\n",
        "                                  label_map=dictlabels)\n",
        "    print(\"Modèle importé\")\n",
        "\n",
        "  else : \n",
        "    model = lp.Detectron2LayoutModel(config_path= config,\n",
        "                                      model_path=modele)\n",
        "                                  #label_map=labels)\n",
        "    print(\"Modèle importé\")\n",
        "except : print(\"Le modèle n'a pas été importé\")"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Modèle importé\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "6u-uS6V3eKdU"
      },
      "source": [
        "#@markdown ## Visualisation de l'inférence sur une page de document\n",
        "\n",
        "\n",
        "#@markdown #### Entrer le chemin vers la page d'un document au format jpg :\n",
        "page = \"\" #@param {type:\"string\"}\n",
        "#@markdown Exemple de chemin: /content/drive/MyDrive/Datasets/Journal/page02.jpg\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown #### Filtrer les résultats :\n",
        "\n",
        "seuil_score = 0.3 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown 0 ou vide pour ne pas filtrer\n",
        "\n",
        "image = cv2.imread(page)\n",
        "image = image[..., ::-1]\n",
        "\n",
        "layout = model.detect(image)\n",
        "if seuil_score :\n",
        "  layout = [x for x in layout if x.score > seuil_score]\n",
        "  \n",
        "for x in layout :\n",
        "  print(x.type,x.score)\n",
        "\n",
        "lp.draw_box(image, layout, box_width=3, show_element_id=True)\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}