mmdet_RTDT-948_Label_studio_train_a_model_on_the_restricted_class (#10)

* added training on restricted class * change object-config.txt
logivations · Jul 12, 2024 · b8f9447 · b8f9447
1 parent c9a78e4
commit b8f9447
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 31 deletions.
diff --git a/auto_training/config_factories/pvt_config_factory.py b/auto_training/config_factories/pvt_config_factory.py
@@ -76,10 +76,8 @@ def make_pvt_cfg(data_path: str,
             times=8,
             dataset=dict(
                 type=cfg.dataset_type,
-                ann_file=
-                cfg.train_ann_file,
-                img_prefix=
-                cfg.train_img_prefix,
+                ann_file=cfg.train_ann_file,
+                img_prefix=cfg.train_img_prefix,
                 pipeline=[
                     dict(type='LoadImageFromFile'),
                     dict(type='LoadAnnotations', with_bbox=True),

diff --git a/auto_training/prepare.py b/auto_training/prepare.py
@@ -1,6 +1,7 @@
 import argparse
 import json
 import os
+from mmcv import DictAction
 
 from auto_training.utils.kitti_conversion import convert_kitti_files
 from auto_training.utils.utils import copy_images
@@ -25,8 +26,7 @@ def make_coco_folder(cocos, coco_path, train_image_path, val_image_path, test_im
 
 def main():
     args = parse_args()
-    target_class_map = json.loads(args.target_class_map)
-    cocos = convert_kitti_files(args.kitti_train, args.kitti_val, args.kitti_test, target_class_map)
+    cocos = convert_kitti_files(args.kitti_train, args.kitti_val, args.kitti_test, args.target_class_map)
     make_coco_folder(cocos, args.coco_folder, args.kitti_train, args.kitti_val, args.kitti_test)
 
 
@@ -36,7 +36,12 @@ def parse_args():
     parser.add_argument('kitti_val', type=str, help='val input data path, kitti dataset')
     parser.add_argument('kitti_test', type=str, help='test input data path, kitti dataset')
     parser.add_argument('coco_folder', type=str, help='output folder')
-    parser.add_argument('--target-class-map', type=str, default="{}", help='target class mapping, json strin format. Map to None if class should not be used.')
+    parser.add_argument(
+        '--target-class-map',
+        nargs='+',
+        action=DictAction,
+        default={},
+        help='target class mapping, json string format. Map to None if class should not be used.')
     return parser.parse_args()
 
 

diff --git a/auto_training/utils/export_specs.py b/auto_training/utils/export_specs.py
@@ -116,9 +116,9 @@ def export_for_lv(args):
     cfg.author = args.author
     export_folder = os.path.join(cfg.work_dir, "export")
     recreate_dir(export_folder)
-    model_name = f"pvt_detector_{cfg.project_name}_{time.strftime('%y%m%d')}"
+    model_name = "pvt_detector"
     write_detector_yaml(cfg=cfg, write_dir=export_folder, name=model_name)
-    write_deepstream_config(cfg=cfg, write_dir=export_folder, name=model_name)
+    #write_deepstream_config(cfg=cfg, write_dir=export_folder, name=model_name)
     write_info_file(cfg=cfg, write_dir=export_folder)
     copy_training_specs(cfg=cfg, write_dir=export_folder)
     print(f"Training info exported successfully to: {export_folder}")

diff --git a/mmdet/datasets/auto.py b/mmdet/datasets/auto.py
@@ -21,28 +21,6 @@
 
 @DATASETS.register_module()
 class AutoDataset(CustomDataset):
-    #CLASSES = ("forklift", "person", "manual_pallet_jack")
-    PALETTE = [(220, 20, 60), (119, 11, 32), (0, 0, 142), (0, 0, 230),
-               (106, 0, 228), (0, 60, 100), (0, 80, 100), (0, 0, 70),
-               (0, 0, 192), (250, 170, 30), (100, 170, 30), (220, 220, 0),
-               (175, 116, 175), (250, 0, 30), (165, 42, 42), (255, 77, 255),
-               (0, 226, 252), (182, 182, 255), (0, 82, 0), (120, 166, 157),
-               (110, 76, 0), (174, 57, 255), (199, 100, 0), (72, 0, 118),
-               (255, 179, 240), (0, 125, 92), (209, 0, 151), (188, 208, 182),
-               (0, 220, 176), (255, 99, 164), (92, 0, 73), (133, 129, 255),
-               (78, 180, 255), (0, 228, 0), (174, 255, 243), (45, 89, 255),
-               (134, 134, 103), (145, 148, 174), (255, 208, 186),
-               (197, 226, 255), (171, 134, 1), (109, 63, 54), (207, 138, 255),
-               (151, 0, 95), (9, 80, 61), (84, 105, 51), (74, 65, 105),
-               (166, 196, 102), (208, 195, 210), (255, 109, 65), (0, 143, 149),
-               (179, 0, 194), (209, 99, 106), (5, 121, 0), (227, 255, 205),
-               (147, 186, 208), (153, 69, 1), (3, 95, 161), (163, 255, 0),
-               (119, 0, 170), (0, 182, 199), (0, 165, 120), (183, 130, 88),
-               (95, 32, 0), (130, 114, 135), (110, 129, 133), (166, 74, 118),
-               (219, 142, 185), (79, 210, 114), (178, 90, 62), (65, 70, 15),
-               (127, 167, 115), (59, 105, 106), (142, 108, 45), (196, 172, 0),
-               (95, 54, 80), (128, 76, 255), (201, 57, 1), (246, 0, 122),
-               (191, 162, 208)]
 
     def load_annotations(self, ann_file):
         """Load annotation from COCO style annotation file.

diff --git a/run_make_nvinfer.sh b/run_make_nvinfer.sh
@@ -0,0 +1,37 @@
+#!/bin/bash
+
+MODEL_DIR=$1
+MODEL_DIR=${MODEL_DIR%%/}
+shift
+
+CLASSES=( "$@" )
+CLASSES=$(IFS=';' ; echo "${CLASSES[*]}")
+
+echo "[property]
+gpu-id=0
+
+# preprocessing parameters.
+net-scale-factor=0.01735207357279195
+offsets=123.675;116.28;103.53
+model-color-format=0
+scaling-filter=3 # 0=Nearest, 1=Bilinear 2=VIC-5 Tap interpolation 3=VIC-10 Tap interpolation
+
+onnx-file=pvt_detector.onnx
+model-engine-file=pvt_detector.onnx_b1_gpu0_fp16.engine
+
+# model config
+infer-dims=3;384;512
+batch-size=1
+network-mode=2 # 0=FP32, 1=INT8, 2=FP16
+network-type=100 # >3 disables post-processing
+cluster-mode=4 # 1=DBSCAN 4=No Clustering
+gie-unique-id=1
+output-tensor-meta=1
+
+[custom]
+min_confidence = 0.5
+labels=$CLASSES
+report_labels=$CLASSES
+" > "$MODEL_DIR/object-config.txt"
+
+