From 2ceacc5e6c2e7b91a4a9b066385544321dc97222 Mon Sep 17 00:00:00 2001
From: Jernej Sabadin <jernej14sabadin@gmail.com>
Date: Thu, 16 Jan 2025 13:47:31 +0100
Subject: [PATCH] small fixes

---
 luxonis_train/config/predefined_models/README.md       |  2 ++
 .../config/predefined_models/detection_fomo_model.py   | 10 +++++-----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/luxonis_train/config/predefined_models/README.md b/luxonis_train/config/predefined_models/README.md
index 629b17a5..a2d29730 100644
--- a/luxonis_train/config/predefined_models/README.md
+++ b/luxonis_train/config/predefined_models/README.md
@@ -162,6 +162,8 @@ See an example configuration file using this predefined model [here](../../../co
 
 There is a trade-off in this simple model: training with a larger `object_weight` in the loss parameters may result in more false positives (FP), but it will improve accuracy. You can also use `use_nms: True` in the `head_params` to enable NMS which can reduce FP, but it will also reduce TP for close neighbors.
 
+For larger heatmaps and improved accuracy, you can adjust the `attach_index` in the `head_params` to a lower value. This will connect the head to an earlier layer in the backbone, resulting in larger heatmaps. However, be aware that this may lead to slower inference times.
+
 ### **Components**
 
 | Name                                                                                            | Alias               | Function                                                                       |
diff --git a/luxonis_train/config/predefined_models/detection_fomo_model.py b/luxonis_train/config/predefined_models/detection_fomo_model.py
index 79ad92aa..2e893a45 100644
--- a/luxonis_train/config/predefined_models/detection_fomo_model.py
+++ b/luxonis_train/config/predefined_models/detection_fomo_model.py
@@ -25,16 +25,16 @@ class FOMOVariant(BaseModel):
 def get_variant(variant: VariantLiteral) -> FOMOVariant:
     """Returns the specific variant configuration for the FOMOModel."""
     variants = {
-        "heavy": FOMOVariant(
-            backbone="MobileNetV2",
-            head_params={"num_conv_layers": 2, "conv_channels": 16},
-            backbone_params={},
-        ),
         "light": FOMOVariant(
             backbone="EfficientRep",
             head_params={"num_conv_layers": 2, "conv_channels": 16},
             backbone_params={"variant": "n"},
         ),
+        "heavy": FOMOVariant(
+            backbone="MobileNetV2",
+            head_params={"num_conv_layers": 2, "conv_channels": 16},
+            backbone_params={},
+        ),
     }
 
     if variant not in variants: