commit with training and evaluation code

.gitignore

@@ -14,8 +14,6 @@ backup/
 
 temp.jpg
 
-Evaluation.ipynb
-
 study/
 
 result_bpmn.bpmn
@@ -24,8 +22,8 @@ BPMN_creation.ipynb
 
 *.png
 
-*.ipynb
-
 *.pmw
 
 best_models.txt
+
+Wizard_creation.ipynb

Evaluation_colab.ipynb

@@ -0,0 +1,758 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "N7fMlFb-n3dJ",
+        "outputId": "ed9bb8ea-42a4-4e07-fdfa-d5a9eca0253f"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Requirement already satisfied: yamlu in /usr/local/lib/python3.10/dist-packages (0.0.17)\n",
+            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.10/dist-packages (from yamlu) (3.7.1)\n",
+            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from yamlu) (1.26.4)\n",
+            "Requirement already satisfied: Pillow in /usr/local/lib/python3.10/dist-packages (from yamlu) (9.4.0)\n",
+            "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->yamlu) (1.2.1)\n",
+            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib->yamlu) (0.12.1)\n",
+            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->yamlu) (4.53.1)\n",
+            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->yamlu) (1.4.5)\n",
+            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->yamlu) (24.1)\n",
+            "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->yamlu) (3.1.2)\n",
+            "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib->yamlu) (2.8.2)\n",
+            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.7->matplotlib->yamlu) (1.16.0)\n",
+            "Requirement already satisfied: optuna in /usr/local/lib/python3.10/dist-packages (3.6.1)\n",
+            "Requirement already satisfied: alembic>=1.5.0 in /usr/local/lib/python3.10/dist-packages (from optuna) (1.13.2)\n",
+            "Requirement already satisfied: colorlog in /usr/local/lib/python3.10/dist-packages (from optuna) (6.8.2)\n",
+            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from optuna) (1.26.4)\n",
+            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from optuna) (24.1)\n",
+            "Requirement already satisfied: sqlalchemy>=1.3.0 in /usr/local/lib/python3.10/dist-packages (from optuna) (2.0.32)\n",
+            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from optuna) (4.66.5)\n",
+            "Requirement already satisfied: PyYAML in /usr/local/lib/python3.10/dist-packages (from optuna) (6.0.2)\n",
+            "Requirement already satisfied: Mako in /usr/local/lib/python3.10/dist-packages (from alembic>=1.5.0->optuna) (1.3.5)\n",
+            "Requirement already satisfied: typing-extensions>=4 in /usr/local/lib/python3.10/dist-packages (from alembic>=1.5.0->optuna) (4.12.2)\n",
+            "Requirement already satisfied: greenlet!=0.4.17 in /usr/local/lib/python3.10/dist-packages (from sqlalchemy>=1.3.0->optuna) (3.0.3)\n",
+            "Requirement already satisfied: MarkupSafe>=0.9.2 in /usr/local/lib/python3.10/dist-packages (from Mako->alembic>=1.5.0->optuna) (2.1.5)\n",
+            "Requirement already satisfied: streamlit in /usr/local/lib/python3.10/dist-packages (1.37.1)\n",
+            "Requirement already satisfied: altair<6,>=4.0 in /usr/local/lib/python3.10/dist-packages (from streamlit) (4.2.2)\n",
+            "Requirement already satisfied: blinker<2,>=1.0.0 in /usr/lib/python3/dist-packages (from streamlit) (1.4)\n",
+            "Requirement already satisfied: cachetools<6,>=4.0 in /usr/local/lib/python3.10/dist-packages (from streamlit) (5.5.0)\n",
+            "Requirement already satisfied: click<9,>=7.0 in /usr/local/lib/python3.10/dist-packages (from streamlit) (8.1.7)\n",
+            "Requirement already satisfied: numpy<3,>=1.20 in /usr/local/lib/python3.10/dist-packages (from streamlit) (1.26.4)\n",
+            "Requirement already satisfied: packaging<25,>=20 in /usr/local/lib/python3.10/dist-packages (from streamlit) (24.1)\n",
+            "Requirement already satisfied: pandas<3,>=1.3.0 in /usr/local/lib/python3.10/dist-packages (from streamlit) (2.1.4)\n",
+            "Requirement already satisfied: pillow<11,>=7.1.0 in /usr/local/lib/python3.10/dist-packages (from streamlit) (9.4.0)\n",
+            "Requirement already satisfied: protobuf<6,>=3.20 in /usr/local/lib/python3.10/dist-packages (from streamlit) (3.20.3)\n",
+            "Requirement already satisfied: pyarrow>=7.0 in /usr/local/lib/python3.10/dist-packages (from streamlit) (14.0.2)\n",
+            "Requirement already satisfied: requests<3,>=2.27 in /usr/local/lib/python3.10/dist-packages (from streamlit) (2.32.3)\n",
+            "Requirement already satisfied: rich<14,>=10.14.0 in /usr/local/lib/python3.10/dist-packages (from streamlit) (13.7.1)\n",
+            "Requirement already satisfied: tenacity<9,>=8.1.0 in /usr/local/lib/python3.10/dist-packages (from streamlit) (8.5.0)\n",
+            "Requirement already satisfied: toml<2,>=0.10.1 in /usr/local/lib/python3.10/dist-packages (from streamlit) (0.10.2)\n",
+            "Requirement already satisfied: typing-extensions<5,>=4.3.0 in /usr/local/lib/python3.10/dist-packages (from streamlit) (4.12.2)\n",
+            "Requirement already satisfied: gitpython!=3.1.19,<4,>=3.0.7 in /usr/local/lib/python3.10/dist-packages (from streamlit) (3.1.43)\n",
+            "Requirement already satisfied: pydeck<1,>=0.8.0b4 in /usr/local/lib/python3.10/dist-packages (from streamlit) (0.9.1)\n",
+            "Requirement already satisfied: tornado<7,>=6.0.3 in /usr/local/lib/python3.10/dist-packages (from streamlit) (6.3.3)\n",
+            "Requirement already satisfied: watchdog<5,>=2.1.5 in /usr/local/lib/python3.10/dist-packages (from streamlit) (4.0.2)\n",
+            "Requirement already satisfied: entrypoints in /usr/local/lib/python3.10/dist-packages (from altair<6,>=4.0->streamlit) (0.4)\n",
+            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from altair<6,>=4.0->streamlit) (3.1.4)\n",
+            "Requirement already satisfied: jsonschema>=3.0 in /usr/local/lib/python3.10/dist-packages (from altair<6,>=4.0->streamlit) (4.23.0)\n",
+            "Requirement already satisfied: toolz in /usr/local/lib/python3.10/dist-packages (from altair<6,>=4.0->streamlit) (0.12.1)\n",
+            "Requirement already satisfied: gitdb<5,>=4.0.1 in /usr/local/lib/python3.10/dist-packages (from gitpython!=3.1.19,<4,>=3.0.7->streamlit) (4.0.11)\n",
+            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.10/dist-packages (from pandas<3,>=1.3.0->streamlit) (2.8.2)\n",
+            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas<3,>=1.3.0->streamlit) (2024.1)\n",
+            "Requirement already satisfied: tzdata>=2022.1 in /usr/local/lib/python3.10/dist-packages (from pandas<3,>=1.3.0->streamlit) (2024.1)\n",
+            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.27->streamlit) (3.3.2)\n",
+            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.27->streamlit) (3.7)\n",
+            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.27->streamlit) (2.0.7)\n",
+            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.27->streamlit) (2024.7.4)\n",
+            "Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.10/dist-packages (from rich<14,>=10.14.0->streamlit) (3.0.0)\n",
+            "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.10/dist-packages (from rich<14,>=10.14.0->streamlit) (2.16.1)\n",
+            "Requirement already satisfied: smmap<6,>=3.0.1 in /usr/local/lib/python3.10/dist-packages (from gitdb<5,>=4.0.1->gitpython!=3.1.19,<4,>=3.0.7->streamlit) (5.0.1)\n",
+            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->altair<6,>=4.0->streamlit) (2.1.5)\n",
+            "Requirement already satisfied: attrs>=22.2.0 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=3.0->altair<6,>=4.0->streamlit) (24.2.0)\n",
+            "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=3.0->altair<6,>=4.0->streamlit) (2023.12.1)\n",
+            "Requirement already satisfied: referencing>=0.28.4 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=3.0->altair<6,>=4.0->streamlit) (0.35.1)\n",
+            "Requirement already satisfied: rpds-py>=0.7.1 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=3.0->altair<6,>=4.0->streamlit) (0.20.0)\n",
+            "Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.10/dist-packages (from markdown-it-py>=2.2.0->rich<14,>=10.14.0->streamlit) (0.1.2)\n",
+            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.2->pandas<3,>=1.3.0->streamlit) (1.16.0)\n",
+            "Mounted at /content/drive\n"
+          ]
+        }
+      ],
+      "source": [
+        "%pip install yamlu\n",
+        "%pip install optuna\n",
+        "%pip install streamlit\n",
+        "\n",
+        "from google.colab import drive\n",
+        "import os\n",
+        "\n",
+        "drive.mount('/content/drive')\n",
+        "path = 'drive/MyDrive/ELCA/BPMN project/'\n",
+        "\n",
+        "os.chdir(path)\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "YkZcbI53n3Dm",
+        "outputId": "18adb94b-567a-46eb-ee84-0b4a77dc00a2"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "100%|██████████| 92/92 [00:30<00:00,  3.04it/s]\n"
+          ]
+        }
+      ],
+      "source": [
+        "from yamlu import ls\n",
+        "from yamlu.coco_read import CocoReader\n",
+        "from pathlib import Path\n",
+        "import cv2\n",
+        "from modules.utils import *\n",
+        "from modules.eval import *\n",
+        "from modules.train import *\n",
+        "from modules.dataset_loader import *\n",
+        "\n",
+        "dataset_path = Path(\"../data/hdBPMN-COCO\")\n",
+        "ls(dataset_path)\n",
+        "\n",
+        "\n",
+        "bpmn_reader = CocoReader(\n",
+        "    dataset_root=dataset_path,\n",
+        "    arrow_categories=[\"sequenceFlow\", \"messageFlow\", \"dataAssociation\"],\n",
+        ")\n",
+        "\n",
+        "\n",
+        "test_anot = bpmn_reader.parse_split(\"test\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Ert1SxZbn3Dn",
+        "outputId": "6384181c-9129-489a-e694-f6b0cd1b57a7"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Loaded 92 annotations.\n"
+          ]
+        }
+      ],
+      "source": [
+        "from torchvision import transforms\n",
+        "from modules.utils import object_dict, arrow_dict, class_dict\n",
+        "from modules.dataset_loader import create_loader\n",
+        "\n",
+        "new_size = (1333,1333)\n",
+        "\n",
+        "model_type = 'object'\n",
+        "\n",
+        "if model_type == 'object':\n",
+        "  model_dict = object_dict\n",
+        "else:\n",
+        "  model_dict = arrow_dict\n",
+        "\n",
+        "transformation_test = transforms.Compose([\n",
+        "    transforms.ToTensor(),\n",
+        "\n",
+        "])\n",
+        "\n",
+        "test_loader = create_loader(new_size, transformation_test, test_anot, batch_size=1, model_type = model_type, seed=42)\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "hp8jehlrXOay"
+      },
+      "outputs": [],
+      "source": [
+        "from modules.train import get_faster_rcnn_model, get_arrow_model\n",
+        "import torch\n",
+        "\n",
+        "# Function to load the models only once and use session state to keep track of it\n",
+        "def load_object_models(model_to_load, model_dict):\n",
+        "  # Adjusted to pass the class_dict directly\n",
+        "  model = get_faster_rcnn_model(len(model_dict))\n",
+        "\n",
+        "  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
+        "  # Load the model weights\n",
+        "  model.load_state_dict(torch.load('./models/'+ model_to_load, map_location=device))\n",
+        "\n",
+        "\n",
+        "  model.to(device)\n",
+        "\n",
+        "  return model\n",
+        "\n",
+        "def load_arrow_models(model_to_load, arrow_dict):\n",
+        "  model = get_arrow_model(len(arrow_dict),2)\n",
+        "\n",
+        "  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
+        "  # Load the model weights\n",
+        "  model.load_state_dict(torch.load('./models/'+ model_to_load, map_location=device))\n",
+        "\n",
+        "\n",
+        "  model.to(device)\n",
+        "\n",
+        "  return model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "6hdMAQ7RX8K8",
+        "outputId": "1315f88b-06c9-42c8-f209-c45abe852e06"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "['model_AdamW_1ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject2.pth', 'model_AdamW_1ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject3.pth', 'model_AdamW_2ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject3.pth', 'model_AdamW_3ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject3.pth', 'model_AdamW_1ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject4.pth', 'model_AdamW_2ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject4.pth', 'model_AdamW_3ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject4.pth', 'model_AdamW_4ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject4.pth', 'model_AdamW_5ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject4.pth', 'model_AdamW_1ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_arrow4.pth', 'model_AdamW_2ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_arrow4.pth', 'model_AdamW_3ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_arrow4.pth', 'model_AdamW_4ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_arrow4.pth', 'model_AdamW_5ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_arrow4.pth']\n",
+            "There is 14 models to test\n"
+          ]
+        }
+      ],
+      "source": [
+        "import os\n",
+        "model_folder = \"models\"\n",
+        "elements = os.listdir(model_folder)\n",
+        "elements = [element for element in elements if \"Adam\" in element]\n",
+        "#elements = [element for element in elements if \"recall\" in element]\n",
+        "print(elements)\n",
+        "print(f\"There is {len(elements)} models to test\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fwz0QKdxgBJz"
+      },
+      "outputs": [],
+      "source": [
+        "from modules.eval import main_evaluation"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "bHkTL_5Jq_t0",
+        "outputId": "026ada96-c865-4a88-c212-8b455d659859"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "There is 8 models to test\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Testing... : 100%|██████████| 92/92 [00:14<00:00,  6.30it/s]\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "1: model_AdamW_1ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject2.pth\n",
+            "Labels_precision: 0.9683, Precision: 0.9742, Recall: 0.9438, F1 Score: 0.9588 \n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Testing... : 100%|██████████| 92/92 [00:14<00:00,  6.38it/s]\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "2: model_AdamW_1ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject3.pth\n",
+            "Labels_precision: 0.9701, Precision: 0.9541, Recall: 0.9600, F1 Score: 0.9571 \n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Testing... : 100%|██████████| 92/92 [00:14<00:00,  6.27it/s]\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "3: model_AdamW_2ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject3.pth\n",
+            "Labels_precision: 0.9701, Precision: 0.9541, Recall: 0.9600, F1 Score: 0.9571 \n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Testing... : 100%|██████████| 92/92 [00:14<00:00,  6.43it/s]\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "4: model_AdamW_3ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject3.pth\n",
+            "Labels_precision: 0.9699, Precision: 0.9658, Recall: 0.9532, F1 Score: 0.9595 \n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Testing... : 100%|██████████| 92/92 [00:14<00:00,  6.38it/s]\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "5: model_AdamW_1ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject4.pth\n",
+            "Labels_precision: 0.9649, Precision: 0.9565, Recall: 0.9607, F1 Score: 0.9586 \n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Testing... : 100%|██████████| 92/92 [00:14<00:00,  6.41it/s]\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "6: model_AdamW_2ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject4.pth\n",
+            "Labels_precision: 0.9704, Precision: 0.9700, Recall: 0.9482, F1 Score: 0.9590 \n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Testing... : 100%|██████���███| 92/92 [00:14<00:00,  6.47it/s]\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "7: model_AdamW_3ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject4.pth\n",
+            "Labels_precision: 0.9708, Precision: 0.9631, Recall: 0.9619, F1 Score: 0.9625 \n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Testing... : 100%|██████████| 92/92 [00:14<00:00,  6.41it/s]"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "8: model_AdamW_4ep_4batch_trainval_blur00_crop02_flip02_rotate02_finetune_bestobject4.pth\n",
+            "Labels_precision: 0.9708, Precision: 0.9631, Recall: 0.9619, F1 Score: 0.9625 \n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "\n"
+          ]
+        }
+      ],
+      "source": [
+        "results = {}\n",
+        "print(f\"There is {len(elements)} models to test\")\n",
+        "for idx, model_name in enumerate(elements):\n",
+        "    if model_type == 'object':\n",
+        "      model = load_object_models(model_name, model_dict)\n",
+        "    else:\n",
+        "      model = load_arrow_models(model_name, model_dict)\n",
+        "\n",
+        "    labels_precision, precision, recall, f1_score, key_accuracy, reverted_accuracy = main_evaluation(model, test_loader,score_threshold=0.5, iou_threshold=0.5, distance_threshold=10, key_correction=False, model_type=model_type)\n",
+        "    print(f\"{idx+1}: {model_name}\")\n",
+        "    print(f\"Labels_precision: {labels_precision:.4f}, Precision: {precision:.4f}, Recall: {recall:.4f}, F1 Score: {f1_score:.4f} \")\n",
+        "    results[model_name] = [labels_precision, precision, recall, f1_score, key_accuracy, reverted_accuracy]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 88
+        },
+        "id": "v0pe9A7DnbUV",
+        "outputId": "80386876-8f54-4166-cb7a-bb7c63cf9414"
+      },
+      "outputs": [
+        {
+          "data": {
+            "application/vnd.google.colaboratory.intrinsic+json": {
+              "type": "string"
+            },
+            "text/plain": [
+              "'for i, metric in enumerate([\\'labels_precision\\', \\'precision\\', \\'recall\\', \\'f1_score\\',\\'key_accuracy\\']):\\n    best_model = max(results, key=lambda x: results[x][i])\\n    print(f\"Best model for {metric}: {best_model}\")\\n    #print all score for this one\\n    print(f\\'Labels Precision: {results[best_model][0]:.3f}, Precision: {results[best_model][1]:.3f}, Recall: {results[best_model][2]:.3f}, F1 Score: {results[best_model][3]:.3f}, Key Accuracy: {results[best_model][4]:.3f}\\')'"
+            ]
+          },
+          "execution_count": 9,
+          "metadata": {},
+          "output_type": "execute_result"
+        }
+      ],
+      "source": [
+        "\"\"\"for i, metric in enumerate(['labels_precision', 'precision', 'recall', 'f1_score','key_accuracy']):\n",
+        "    best_model = max(results, key=lambda x: results[x][i])\n",
+        "    print(f\"Best model for {metric}: {best_model}\")\n",
+        "    #print all score for this one\n",
+        "    print(f'Labels Precision: {results[best_model][0]:.3f}, Precision: {results[best_model][1]:.3f}, Recall: {results[best_model][2]:.3f}, F1 Score: {results[best_model][3]:.3f}, Key Accuracy: {results[best_model][4]:.3f}')\"\"\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "HMyYdPjLiGMH",
+        "outputId": "b5e28040-9703-4d3c-9aeb-8ab945a78c21"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Downloading: \"https://download.pytorch.org/models/resnet50-0676ba61.pth\" to /root/.cache/torch/hub/checkpoints/resnet50-0676ba61.pth\n",
+            "100%|██████████| 97.8M/97.8M [00:03<00:00, 31.2MB/s]\n",
+            "Testing... : 100%|██████████| 92/92 [00:20<00:00,  4.44it/s]"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "best_model_object.pth\n",
+            "Labels_precision: 0.9671, Precision: 0.9429, Recall: 0.9682, F1 Score: 0.9553\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "\n"
+          ]
+        }
+      ],
+      "source": [
+        "from modules.eval import main_evaluation\n",
+        "\n",
+        "\n",
+        "results = {}\n",
+        "model_name = 'best_model_object.pth'\n",
+        "model_dict = object_dict\n",
+        "model = load_object_models(model_name, model_dict)\n",
+        "\n",
+        "labels_precision, precision, recall, f1_score, key_accuracy, reverted_accuracy = main_evaluation(model, test_loader,score_threshold=0.5, iou_threshold=0.5, model_type=model_type)\n",
+        "print(model_name)\n",
+        "print(f\"Labels_precision: {labels_precision:.4f}, Precision: {precision:.4f}, Recall: {recall:.4f}, F1 Score: {f1_score:.4f}\")\n",
+        "#results[model_name] = [labels_precision, precision, recall, f1_score, key_accuracy, reverted_accuracy]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "r6yDD7CljRXA",
+        "outputId": "53eb3edc-7dfd-47fc-e9f8-72b208aefd6e"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Testing... : 100%|██████████| 92/92 [00:15<00:00,  5.83it/s]"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "Class Precision: {'background': 0, 'task': 0.967741935483871, 'exclusiveGateway': 0.9433962264150944, 'event': 0.9461077844311377, 'parallelGateway': 0.926829268292683, 'messageEvent': 0.9230769230769231, 'pool': 0.7453416149068323, 'lane': 0.8554216867469879, 'dataObject': 0.8651685393258427, 'dataStore': 1.0, 'subProcess': 0.0, 'eventBasedGateway': 0.7272727272727273, 'timerEvent': 0.7916666666666666}\n",
+            "Class Recall: {'background': 0, 'task': 0.9810671256454389, 'exclusiveGateway': 0.9554140127388535, 'event': 0.9294117647058824, 'parallelGateway': 0.9344262295081968, 'messageEvent': 0.9523809523809523, 'pool': 0.96, 'lane': 0.71, 'dataObject': 0.9565217391304348, 'dataStore': 0.64, 'subProcess': 0, 'eventBasedGateway': 0.7272727272727273, 'timerEvent': 0.7916666666666666}\n",
+            "Class F1 Score: {'background': 0, 'task': 0.9743589743589743, 'exclusiveGateway': 0.949367088607595, 'event': 0.9376854599406529, 'parallelGateway': 0.9306122448979592, 'messageEvent': 0.9375, 'pool': 0.8391608391608391, 'lane': 0.7759562841530054, 'dataObject': 0.9085545722713865, 'dataStore': 0.7804878048780487, 'subProcess': 0, 'eventBasedGateway': 0.7272727272727273, 'timerEvent': 0.7916666666666666}\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "\n"
+          ]
+        }
+      ],
+      "source": [
+        "class_precision, class_recall, class_f1_score = evaluate_model_by_class(model, test_loader, model_dict, score_threshold=0.5, iou_threshold=0.5)\n",
+        "print(f\"\\nClass Precision: {class_precision}\")\n",
+        "print(f\"Class Recall: {class_recall}\")\n",
+        "print(f\"Class F1 Score: {class_f1_score}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "1wtvRs4zqoDN",
+        "outputId": "08b8f742-2ef3-4414-d84f-e9d089d32b16"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Average Precision: 0.9429\n",
+            "Average Recall: 0.9682\n",
+            "Average F1 Score: 0.9553\n"
+          ]
+        }
+      ],
+      "source": [
+        "import numpy as np\n",
+        "\n",
+        "#average each\n",
+        "average_precision = np.mean(precision)\n",
+        "average_recall = np.mean(recall)\n",
+        "average_f1_score = np.mean(f1_score)\n",
+        "\n",
+        "print(f\"Average Precision: {average_precision:.4f}\")\n",
+        "print(f\"Average Recall: {average_recall:.4f}\")\n",
+        "print(f\"Average F1 Score: {average_f1_score:.4f}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "aHVvDOEvKdL4",
+        "outputId": "f6e636aa-d281-4e67-de43-f1783c06194b"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Loaded 92 annotations.\n"
+          ]
+        }
+      ],
+      "source": [
+        "from torchvision import transforms\n",
+        "#from modules.utils import object_dict, arrow_dict, class_dict\n",
+        "\n",
+        "#new_size = (640, 384)\n",
+        "new_size = (1333,1333)\n",
+        "\n",
+        "model_type = 'arrow'\n",
+        "\n",
+        "if model_type == 'object':\n",
+        "  model_dict = object_dict\n",
+        "else:\n",
+        "  model_dict = arrow_dict\n",
+        "\n",
+        "transformation_test = transforms.Compose([\n",
+        "    transforms.ToTensor(),\n",
+        "\n",
+        "])\n",
+        "\n",
+        "test_loader = create_loader(new_size, transformation_test, test_anot, batch_size=1, model_type = model_type)\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "gIyJdC3shmGU",
+        "outputId": "eccaf29a-b01a-460c-fada-34fbd3f626bf"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Testing... : 100%|██████████| 92/92 [00:19<00:00,  4.69it/s]"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            " best_model_arrow.pth\n",
+            "Labels_precision: 0.9873, Precision: 0.9203, Recall: 0.9256, F1 Score: 0.9229, Key Accuracy: 0.7065, Reverted Accuracy: 0.0196\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "\n"
+          ]
+        }
+      ],
+      "source": [
+        "from modules.eval import main_evaluation\n",
+        "\n",
+        "results = {}\n",
+        "model_name = 'best_model_arrow.pth'\n",
+        "model = load_arrow_models(model_name, model_dict)\n",
+        "\n",
+        "for i in range(5):\n",
+        "  test_loader = create_loader(new_size, transformation_test, test_anot, batch_size=1, model_type = model_type, seed=42+i)\n",
+        "  labels_precision, precision, recall, f1_score, key_accuracy, reverted_accuracy = main_evaluation(model, test_loader,score_threshold=0.7, iou_threshold=0.5, distance_threshold=10, key_correction=False, model_type=model_type)\n",
+        "  print(\"\\n\",model_name)\n",
+        "  print(f\"Seed: {42+i} ,Labels_precision: {labels_precision:.4f}, Precision: {precision:.4f}, Recall: {recall:.4f}, F1 Score: {f1_score:.4f}, Key Accuracy: {key_accuracy:.4f}, Reverted Accuracy: {reverted_accuracy:.4f}\")\n",
+        "  #results[model_name] = [labels_precision, precision, recall, f1_score, key_accuracy, reverted_accuracy]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "KIUAasG5hzw1",
+        "outputId": "6a5617b1-ed1b-4237-dee2-3fa40f14f99a"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Testing... : 100%|██████████| 92/92 [00:19<00:00,  4.78it/s]"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Class Precision: {'background': 0, 'sequenceFlow': 0.9075697211155378, 'dataAssociation': 0.7788778877887789, 'messageFlow': 0.7914110429447853}\n",
+            "Class Recall: {'background': 0, 'sequenceFlow': 0.9366776315789473, 'dataAssociation': 0.7492063492063492, 'messageFlow': 0.7288135593220338}\n",
+            "Class F1 Score: {'background': 0, 'sequenceFlow': 0.9218939700526103, 'dataAssociation': 0.7637540453074433, 'messageFlow': 0.7588235294117648}\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "\n"
+          ]
+        }
+      ],
+      "source": [
+        "from modules.eval import evaluate_model_by_class\n",
+        "\n",
+        "class_precision, class_recall, class_f1_score = evaluate_model_by_class(model, test_loader, model_dict, score_threshold=0.7, iou_threshold=0.6)\n",
+        "print(f\"Class Precision: {class_precision}\")\n",
+        "print(f\"Class Recall: {class_recall}\")\n",
+        "print(f\"Class F1 Score: {class_f1_score}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fwkbOQ8Yq019"
+      },
+      "outputs": [],
+      "source": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "machine_shape": "hm",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.12.2"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

modules/train.py

@@ -87,7 +87,7 @@ def prepare_model(dict, opti, learning_rate=0.0003, model_to_load=None, model_ty
     device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
     # Load the model weights
     if model_to_load:
-        model.load_state_dict(torch.load('./models/' + model_to_load + '.pth', map_location=device))
+        model.load_state_dict(torch.load(model_to_load + '.pth', map_location=device))
         print(f"Model '{model_to_load}'  loaded")
 
     model.to(device)
@@ -191,228 +191,187 @@ def evaluate_loss(model, data_loader, device, loss_config=None, print_losses=Fal
 
 
 def training_model(num_epochs, model, data_loader, subset_test_loader,
-                   optimizer, model_to_load=None, change_learning_rate=100, start_key=100,
+                   optimizer, model_to_load=None, change_learning_rate=100, start_key=100, save_every=5,
                    parameters=None, blur_prob=0.02,
                    score_threshold=0.7, iou_threshold=0.5, early_stop_f1_score=0.97,
                    information_training='training', start_epoch=0, loss_config=None, model_type='object',
                    eval_metric='f1_score', device=torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')):
-    """
-    Train the model over a specified number of epochs.
-
-    Parameters:
-    - num_epochs (int): Number of epochs to train for.
-    - model (torch.nn.Module): Model to train.
-    - data_loader (torch.utils.data.DataLoader): DataLoader for the training dataset.
-    - subset_test_loader (torch.utils.data.DataLoader): DataLoader for the validation dataset.
-    - optimizer (torch.optim.Optimizer): Optimizer to use for training.
-    - model_to_load (str, optional): Name of the model to load.
-    - change_learning_rate (int): Epoch interval to change the learning rate.
-    - start_key (int): Epoch to start training keypoints.
-    - parameters (dict, optional): Additional training parameters.
-    - blur_prob (float): Probability of applying blur augmentation.
-    - score_threshold (float): Score threshold for evaluation.
-    - iou_threshold (float): IoU threshold for evaluation.
-    - early_stop_f1_score (float): F1 score threshold for early stopping.
-    - information_training (str): Information about the training.
-    - start_epoch (int): Starting epoch number.
-    - loss_config (dict, optional): Configuration specifying which losses to use.
-    - model_type (str): Type of model ('object' or 'arrow').
-    - eval_metric (str): Evaluation metric ('f1_score', 'precision', 'recall', or 'loss').
-    - device (torch.device): Device to perform training on.
-
-    Returns:
-    - model (torch.nn.Module): Trained model.
-    """
-    model.train()
-
-    if loss_config is None:
-        print('No loss config found, all losses will be used.')
-    else:
-        # Print the list of the losses that will be used
-        print('The following losses will be used: ', end='')
-        for key, value in loss_config.items():
-            if value:
-                print(key, end=", ")
-        print()
-
-    # Initialize lists to store epoch-wise average losses
-    epoch_avg_losses = []
-    epoch_avg_loss_classifier = []
-    epoch_avg_loss_box_reg = []
-    epoch_avg_loss_objectness = []
-    epoch_avg_loss_rpn_box_reg = []
-    epoch_avg_loss_keypoints = []
-    epoch_precision = []
-    epoch_recall = []
-    epoch_f1_score = []
-    epoch_test_loss = []
-
-    start_tot = time.time()
-    best_metrics = -1000
-    best_epoch = 0
-    best_model_state = None
-    same = 0
-    learning_rate = optimizer.param_groups[0]['lr']
-    bad_test_loss = 0
-    previous_test_loss = 1000
-
-    if parameters is not None:
-        batch_size, crop_prob, rotate_90_proba, h_flip_prob, v_flip_prob, max_rotate_deg, rotate_proba, keep_ratio = parameters.values()
-
-    print(f"Let's go training {model_type} model with {num_epochs} epochs!")
-    if parameters is not None:
-        print(f"Learning rate: {learning_rate}, Batch size: {batch_size}, Crop prob: {crop_prob}, H flip prob: {h_flip_prob}, V flip prob: {v_flip_prob}, Max rotate deg: {max_rotate_deg}, Rotate proba: {rotate_proba}, Rotate 90 proba: {rotate_90_proba}, Keep ratio: {keep_ratio}")
-
-    for epoch in range(num_epochs):
-        if (epoch > 0 and (epoch) % change_learning_rate == 0) or bad_test_loss >= 3:
-            learning_rate = 0.7 * learning_rate
-            optimizer = AdamW(model.parameters(), lr=learning_rate, weight_decay=learning_rate, eps=1e-08, betas=(0.9, 0.999))
-            if best_model_state is not None:
-                model.load_state_dict(best_model_state)
-            print(f'Learning rate changed to {learning_rate:.4} and the best epoch for now is {best_epoch}')
-            bad_test_loss = 0
-        if epoch > 0 and (epoch) == start_key:
-            print("Now it's training Keypoints also")
-            loss_config['loss_keypoint'] = True
-            for name, param in model.named_parameters():
-                if 'keypoint' in name:
-                    param.requires_grad = True
-
-        model.train()
-        start = time.time()
-        total_loss = 0
-
-        # Initialize lists to keep track of individual losses
-        loss_classifier_list = []
-        loss_box_reg_list = []
-        loss_objectness_list = []
-        loss_rpn_box_reg_list = []
-        loss_keypoints_list = []
-
-        # Create a tqdm progress bar
-        progress_bar = tqdm(data_loader, desc=f'Epoch {epoch + 1 + start_epoch}')
-
-        for images, targets_im in progress_bar:
-            images = [image.to(device) for image in images]
-            targets = [{k: v.clone().detach().to(device) for k, v in t.items()} for t in targets_im]
-
-            optimizer.zero_grad()
-
-            loss_dict = model(images, targets)
-            # Inside the training loop where losses are calculated:
-            losses = 0
-            if loss_config is not None:
-                for key, loss in loss_dict.items():
-                    if loss_config.get(key, False):
-                        if key == 'loss_classifier':
-                            loss *= 3
-                        losses += loss
-            else:
-                losses = sum(loss for key, loss in loss_dict.items())
-
-            # Collect individual losses
-            if loss_dict['loss_classifier']:
-                loss_classifier_list.append(loss_dict['loss_classifier'].item())
-            else:
-                loss_classifier_list.append(0)
-
-            if loss_dict['loss_box_reg']:
-                loss_box_reg_list.append(loss_dict['loss_box_reg'].item())
-            else:
-                loss_box_reg_list.append(0)
-
-            if loss_dict['loss_objectness']:
-                loss_objectness_list.append(loss_dict['loss_objectness'].item())
-            else:
-                loss_objectness_list.append(0)
-
-            if loss_dict['loss_rpn_box_reg']:
-                loss_rpn_box_reg_list.append(loss_dict['loss_rpn_box_reg'].item())
-            else:
-                loss_rpn_box_reg_list.append(0)
 
-            if 'loss_keypoint' in loss_dict:
-                loss_keypoints_list.append(loss_dict['loss_keypoint'].item())
-            else:
-                loss_keypoints_list.append(0)
-
-            losses.backward()
-            optimizer.step()
-
-            total_loss += losses.item()
+  # Set the model to training mode
+  model.train()
+
+  if loss_config is None:
+     print('No loss config found, all losses will be used.')
+  else:
+     # Print the list of the losses that will be used
+      print('The following losses will be used: ', end='')
+      for key, value in loss_config.items():
+          if value:
+              print(key, end=", ")
+      print()
+
+  # Initialize lists to store epoch-wise average losses and other metrics
+  epoch_avg_losses = []
+  epoch_avg_loss_classifier = []
+  epoch_avg_loss_box_reg = []
+  epoch_avg_loss_objectness = []
+  epoch_avg_loss_rpn_box_reg = []
+  epoch_avg_loss_keypoints = []
+  epoch_precision = []
+  epoch_recall = []
+  epoch_f1_score = []
+  epoch_test_loss = []
+
+  start_tot = time.time()
+  best_metric_value = -1000
+  best_epoch = 0
+  best_model_state = None
+  epochs_with_high_f1 = 0
+  learning_rate = optimizer.param_groups[0]['lr']
+  bad_test_loss_epochs = 0
+  previous_test_loss = 1000
+
+  if parameters is not None:
+    batch_size, crop_prob, rotate_90_proba, h_flip_prob, v_flip_prob, max_rotate_deg, rotate_proba, keep_ratio = parameters.values()
+
+  print(f"Let's go training {model_type} model with {num_epochs} epochs!")
+  if parameters is not None:
+    print(f"Learning rate: {learning_rate}, Batch size: {batch_size}, Crop prob: {crop_prob}, H flip prob: {h_flip_prob}, V flip prob: {v_flip_prob}, Max rotate deg: {max_rotate_deg}, Rotate proba: {rotate_proba}, Rotate 90 proba: {rotate_90_proba}, Keep ratio: {keep_ratio}")
+
+  for epoch in range(num_epochs):
+
+      if (epoch > 0 and epoch % change_learning_rate == 0) or bad_test_loss_epochs >= 2:
+        learning_rate *= 0.7
+        optimizer = AdamW(model.parameters(), lr=learning_rate, weight_decay=learning_rate, eps=1e-08, betas=(0.9, 0.999))
+        if best_model_state is not None:
+          model.load_state_dict(best_model_state)
+        print(f'Learning rate changed to {learning_rate:.4} and the best epoch for now is {best_epoch}')
+        bad_test_loss_epochs = 0
+
+      if epoch > 0 and epoch == start_key:
+        print("Now it's training Keypoints also")
+        loss_config['loss_keypoint'] = True
+        for name, param in model.named_parameters():
+          if 'keypoint' in name:
+              param.requires_grad = True
+
+      model.train()
+      start = time.time()
+      total_loss = 0
+
+      # Initialize lists to keep track of individual losses
+      loss_classifier_list = []
+      loss_box_reg_list = []
+      loss_objectness_list = []
+      loss_rpn_box_reg_list = []
+      loss_keypoints_list = []
+
+      # Create a tqdm progress bar
+      progress_bar = tqdm(data_loader, desc=f'Epoch {epoch+1+start_epoch}')
+
+      for images, targets_im in progress_bar:
+          images = [image.to(device) for image in images]
+          targets = [{k: v.clone().detach().to(device) for k, v in t.items()} for t in targets_im]
+
+          optimizer.zero_grad()
+
+          loss_dict = model(images, targets)
+          # Inside the training loop where losses are calculated:
+          losses = 0
+          if loss_config is not None:
+            for key, loss in loss_dict.items():
+                if loss_config.get(key, False):
+                    if key == 'loss_classifier':
+                      loss *= 3
+                    losses += loss
+          else:
+            losses = sum(loss for key, loss in loss_dict.items())
+
+          # Collect individual losses
+          loss_classifier_list.append(loss_dict.get('loss_classifier', torch.tensor(0)).item())
+          loss_box_reg_list.append(loss_dict.get('loss_box_reg', torch.tensor(0)).item())
+          loss_objectness_list.append(loss_dict.get('loss_objectness', torch.tensor(0)).item())
+          loss_rpn_box_reg_list.append(loss_dict.get('loss_rpn_box_reg', torch.tensor(0)).item())
+          loss_keypoints_list.append(loss_dict.get('loss_keypoint', torch.tensor(0)).item())
+
+          losses.backward()
+          optimizer.step()
+
+          total_loss += losses.item()
+
+          # Update the description with the current loss
+          progress_bar.set_description(f'Epoch {epoch+1+start_epoch}, Loss: {losses.item():.4f}')
+
+      # Calculate average loss
+      avg_loss = total_loss / len(data_loader)
+
+      epoch_avg_losses.append(avg_loss)
+      epoch_avg_loss_classifier.append(np.mean(loss_classifier_list))
+      epoch_avg_loss_box_reg.append(np.mean(loss_box_reg_list))
+      epoch_avg_loss_objectness.append(np.mean(loss_objectness_list))
+      epoch_avg_loss_rpn_box_reg.append(np.mean(loss_rpn_box_reg_list))
+      epoch_avg_loss_keypoints.append(np.mean(loss_keypoints_list))
+
+      # Evaluate the model on the test set
+      if eval_metric == 'loss':
+        labels_precision, precision, recall, f1_score, key_accuracy, reverted_accuracy = 0, 0, 0, 0, 0, 0
+        avg_test_loss = evaluate_loss(model, subset_test_loader, device, loss_config)
+        print(f"Epoch {epoch+1+start_epoch}, Average Training Loss: {avg_loss:.4f}, Average Test Loss: {avg_test_loss:.4f}", end=", ")
+      else:
+        avg_test_loss = 0
+        labels_precision, precision, recall, f1_score, key_accuracy, reverted_accuracy = main_evaluation(model, subset_test_loader, score_threshold=score_threshold, iou_threshold=iou_threshold, distance_threshold=10, key_correction=False, model_type=model_type)
+        print(f"Epoch {epoch+1+start_epoch}, Average Loss: {avg_loss:.4f}, Labels_precision: {labels_precision:.4f}, Precision: {precision:.4f}, Recall: {recall:.4f}, F1 Score: {f1_score:.4f} ", end=", ")
+        avg_test_loss = evaluate_loss(model, subset_test_loader, device, loss_config)
+        print(f"Epoch {epoch+1+start_epoch}, Average Test Loss: {avg_test_loss:.4f}", end=", ")
+
+      print(f"Time: {time.time() - start:.2f} [s]")
+
+      if eval_metric == 'f1_score':
+        metric_used = f1_score
+      elif eval_metric == 'precision':
+        metric_used = precision
+      elif eval_metric == 'recall':
+        metric_used = recall
+      else:
+        metric_used = -avg_test_loss
+
+      # Check if this epoch's model has the best evaluation metric
+      if metric_used > best_metric_value:
+          best_metric_value = metric_used
+          best_epoch = epoch + 1 + start_epoch
+          best_model_state = copy.deepcopy(model.state_dict())
+
+      if epoch > 0 and f1_score > early_stop_f1_score:
+        epochs_with_high_f1 += 1
+
+      epoch_precision.append(precision)
+      epoch_recall.append(recall)
+      epoch_f1_score.append(f1_score)
+      epoch_test_loss.append(avg_test_loss)
+
+      name_model = f"model_{type(optimizer).__name__}_{epoch+1+start_epoch}ep_{batch_size}batch_trainval_blur0{int(blur_prob*10)}_crop0{int(crop_prob*10)}_flip0{int(h_flip_prob*10)}_rotate0{int(rotate_proba*10)}_{information_training}"
+      metrics_list = [epoch_avg_losses, epoch_avg_loss_classifier, epoch_avg_loss_box_reg, epoch_avg_loss_objectness, epoch_avg_loss_rpn_box_reg, epoch_avg_loss_keypoints, epoch_precision, epoch_recall, epoch_f1_score, epoch_test_loss]
+
+      if epochs_with_high_f1 >= 1:
+        torch.save(best_model_state, './models/' + name_model + '.pth')
+        write_results(name_model, metrics_list, start_epoch)
+        break
 
-            # Update the description with the current loss
-            progress_bar.set_description(f'Epoch {epoch + 1 + start_epoch}, Loss: {losses.item():.4f}')
-
-        # Calculate average loss
-        avg_loss = total_loss / len(data_loader)
-
-        epoch_avg_losses.append(avg_loss)
-        epoch_avg_loss_classifier.append(np.mean(loss_classifier_list))
-        epoch_avg_loss_box_reg.append(np.mean(loss_box_reg_list))
-        epoch_avg_loss_objectness.append(np.mean(loss_objectness_list))
-        epoch_avg_loss_rpn_box_reg.append(np.mean(loss_rpn_box_reg_list))
-        epoch_avg_loss_keypoints.append(np.mean(loss_keypoints_list))
-
-        # Evaluate the model on the test set
-        if eval_metric == 'loss':
-            labels_precision, precision, recall, f1_score, key_accuracy, reverted_accuracy = 0, 0, 0, 0, 0, 0
-            avg_test_loss = evaluate_loss(model, subset_test_loader, device, loss_config)
-            print(f"Epoch {epoch + 1 + start_epoch}, Average Training Loss: {avg_loss:.4f}, Average Test Loss: {avg_test_loss:.4f}", end=", ")
-        else:
-            avg_test_loss = 0
-            labels_precision, precision, recall, f1_score, key_accuracy, reverted_accuracy = main_evaluation(model, subset_test_loader, score_threshold=0.5, iou_threshold=0.5, distance_threshold=10, key_correction=False, model_type=model_type)
-            print(f"Epoch {epoch + 1 + start_epoch}, Average Loss: {avg_loss:.4f}, Labels_precision: {labels_precision:.4f}, Precision: {precision:.4f}, Recall: {recall:.4f}, F1 Score: {f1_score:.4f} ", end=", ")
-            avg_test_loss = evaluate_loss(model, subset_test_loader, device, loss_config)
-            print(f"Epoch {epoch + 1 + start_epoch}, Average Test Loss: {avg_test_loss:.4f}", end=", ")
-
-        print(f"Time: {time.time() - start:.2f} [s]")
-
-        if eval_metric == 'f1_score':
-            metric_used = f1_score
-        elif eval_metric == 'precision':
-            metric_used = precision
-        elif eval_metric == 'recall':
-            metric_used = recall
-        else:
-            metric_used = -avg_test_loss
-
-        # Check if this epoch's model has the lowest average loss
-        if metric_used > best_metrics:
-            best_metrics = metric_used
-            best_epoch = epoch + 1 + start_epoch
-            best_model_state = copy.deepcopy(model.state_dict())
-
-        if epoch > 0 and f1_score > early_stop_f1_score:
-            same += 1
-
-        epoch_precision.append(precision)
-        epoch_recall.append(recall)
-        epoch_f1_score.append(f1_score)
-        epoch_test_loss.append(avg_test_loss)
-
-        name_model = f"model_{type(optimizer).__name__}_{epoch + 1 + start_epoch}ep_{batch_size}batch_trainval_blur0{int(blur_prob * 10)}_crop0{int(crop_prob * 10)}_flip0{int(h_flip_prob * 10)}_rotate0{int(rotate_proba * 10)}_{information_training}"
-        metrics_list = [epoch_avg_losses, epoch_avg_loss_classifier, epoch_avg_loss_box_reg, epoch_avg_loss_objectness, epoch_avg_loss_rpn_box_reg, epoch_avg_loss_keypoints, epoch_precision, epoch_recall, epoch_f1_score, epoch_test_loss]
-
-        if same >= 1:
-            torch.save(best_model_state, './models/' + name_model + '.pth')
-            write_results(name_model, metrics_list, start_epoch)
-            break
-
-        if (epoch + 1 + start_epoch) % 5 == 0:
-            torch.save(best_model_state, './models/' + name_model + '.pth')
-            model.load_state_dict(best_model_state)
-            write_results(name_model, metrics_list, start_epoch)
-
-        if avg_test_loss > previous_test_loss:
-            bad_test_loss += 1
-        previous_test_loss = avg_test_loss
-
-    print(f"\n Total time: {(time.time() - start_tot) / 60} minutes, Best Epoch is {best_epoch} with an {eval_metric} of {best_metrics:.4f}")
-    if best_model_state:
+      if (epoch + 1 + start_epoch) % save_every == 0:
         torch.save(best_model_state, './models/' + name_model + '.pth')
         model.load_state_dict(best_model_state)
         write_results(name_model, metrics_list, start_epoch)
-        print(f"Name of the best model: model_{type(optimizer).__name__}_{epoch + 1 + start_epoch}ep_{batch_size}batch_trainval_blur0{int(blur_prob * 10)}_crop0{int(crop_prob * 10)}_flip0{int(h_flip_prob * 10)}_rotate0{int(rotate_proba * 10)}_{information_training}")
 
-    return model
+      if avg_test_loss > previous_test_loss:
+        bad_test_loss_epochs += 1
+      previous_test_loss = avg_test_loss
+
+  print(f"\nTotal time: {(time.time() - start_tot) / 60:.2f} minutes, Best Epoch is {best_epoch} with an {eval_metric} of {best_metric_value:.4f}")
+
+  if best_model_state:
+      torch.save(best_model_state, './models/' + name_model + '.pth')
+      model.load_state_dict(best_model_state)
+      write_results(name_model, metrics_list, start_epoch)
+      print(f"Name of the best model: {name_model}")
+
+  return model