fix rcnn and ssd

src/data_processing/train_processing/convert_yolo_to_voc.py

@@ -5,7 +5,7 @@
 from loguru import logger
 from PIL import Image
 from tqdm import tqdm
-from prepare_ssd import create_ssd_json
+from src.data_processing.train_processing.prepare_ssd import create_ssd_json
 
 
 
@@ -51,7 +51,7 @@ def is_number(n):
             return False
 
     folder_holding_yolo_files = jpeg_images_folder
-    yolo_class_list_file = f"{current_file_path.parents[2]}/{txt_path}"
+    yolo_class_list_file = f"{current_file_path.parents[3]}/{txt_path}"
 
     # Get a list of all the classes used in the YOLO format
     with open(yolo_class_list_file) as f:

src/data_processing/train_processing/prepare_ssd.py

@@ -3,7 +3,7 @@
 import xml.etree.ElementTree as ET
 from tqdm import tqdm
 from pathlib import Path
-from data_utils.utils import load_class_names
+from src.data_processing.data_utils.utils import load_class_names
 
 
 def check_filename(filename):
@@ -44,16 +44,19 @@ def save_as_json(basename, dataset):
 
 def get_image_names(folder_path):
     image_names = []
+    image_extension = []
     for filename in os.listdir(folder_path):
         if filename.endswith(('.jpg', '.jpeg', '.png', '.gif')):
             name = os.path.splitext(filename)[0]
+            extension = os.path.splitext(filename)[-1]
             image_names.append(name)
-    return image_names
+            image_extension.append(extension)
+    return (image_names, image_extension)
 
 
 def create_ssd_json(path_folder, txt_path):
     current_file_path = Path(__file__).resolve()
-    txt_path = Path(current_file_path.parents[2]) / txt_path
+    txt_path = Path(current_file_path.parents[3]) / txt_path
     class_names = load_class_names(txt_path)
 
     paths = {
@@ -62,9 +65,9 @@ def create_ssd_json(path_folder, txt_path):
 
     dataset = []
     for year, path in paths.items():
-        ids = get_image_names(Path(path_folder) / 'images')
-        for id in tqdm(ids):
-            image_path = os.path.join(path, 'images', id + '.jpg')
+        ids, ids_extentions = get_image_names(Path(path_folder) / 'images')
+        for i, id in enumerate(tqdm(ids)):
+            image_path = os.path.join(path, 'images', id + ids_extentions[i])
             annotation_path = os.path.join(path, 'annotations', id + '.xml')
             if check_filename(annotation_path):
                 try:

src/data_processing/train_processing/prepare_train.py

@@ -15,7 +15,7 @@ def model_selection(MODEL):
     arch = ""
     if MODEL.startswith('yolov5'):
         arch = 'yolov5'
-        path_config = Path(file.parents[1]) / 'train_utils' / 'train_model' / 'models' / 'yolov5' / 'models' / f'{MODEL}.yaml'
+        path_config = Path(file.parents[2]) / 'train_utils' / 'train_models' / 'models' / 'yolov5' / 'models' / f'{MODEL}.yaml'
         if os.path.exists(path_config):
             return arch, path_config
         else:
@@ -25,7 +25,7 @@ def model_selection(MODEL):
     elif MODEL.startswith('yolov7'):
         arch = 'yolov7'
         path_config = (
-            Path(file.parents[1]) / 'train_utils' / 'train_model' / 'models' /
+            Path(file.parents[2]) / 'train_utils' / 'train_models' / 'models' /
             'yolov7' / 'cfg' / 'training' / f'{MODEL}.yaml'
             )
         if os.path.exists(path_config):
@@ -37,7 +37,7 @@ def model_selection(MODEL):
     elif MODEL.startswith('yolov8'):
         arch = 'yolov8'
         path_config = (
-            Path(file.parents[1]) / 'train_utils' / 'train_model' / 'models' /
+            Path(file.parents[2]) / 'train_utils' / 'train_models' / 'models' /
             'ultralytics' / 'ultralytics' / 'models' / 'v8' / f'{MODEL}.yaml'
             )
         if os.path.exists(path_config):
@@ -84,7 +84,7 @@ def create_config_data(train_path, val_path, classname_file, config_path, arch,
     current_file_path = Path(__file__).resolve()
 
     runs_path = create_run_directory(model)
-    class_file_path = Path(current_file_path.parents[2]) / classname_file
+    class_file_path = Path(current_file_path.parents[3]) / classname_file
 
     config_path = runs_path / config_path
     if arch == 'ssd':
@@ -170,3 +170,11 @@ def create_config_data(train_path, val_path, classname_file, config_path, arch,
 
         return config_path
 
+def check_config_arrays_sizes(dictionary):
+    for key, value in dictionary.items():
+        if isinstance(value, list):
+            first_array = next(iter(dictionary.values()))
+            first_array_size = len(first_array)
+            current_array_size = len(value)
+            if current_array_size != first_array_size:
+                raise ValueError(f"Size mismatch for key '{key}'. Expected size: {first_array_size}, actual size: {current_array_size}")

src/train_utils/config/train_config.yaml

@@ -1,12 +1,12 @@
 BATCH_SIZE: 20
 CLASSES: classes.txt
-DATA_PATH: /home/runner/work/ODRS/ODRS/user_datasets/WaRP/Warp-D
-EPOCHS: 1
+#DATA_PATH: /home/runner/work/ODRS/ODRS/user_datasets/WaRP/Warp-D
+DATA_PATH: mon_part_1
+EPOCHS: 4
 GPU_COUNT: 1
 IMG_SIZE: 300
-MODEL: yolov8s
-SELECT_GPU: cpu
+MODEL: ssd
+SELECT_GPU: 0
 
-CONFIG_PATH: dataset.yaml
 SPLIT_TRAIN_VALUE: 0.6
 SPLIT_VAL_VALUE: 0.35

src/train_utils/train_detectors.py

@@ -4,16 +4,16 @@
 from loguru import logger
 project_dir = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(os.path.dirname(os.path.dirname(project_dir)))
-from src.data_processing.split_dataset import split_data, copy_arch_folder
-from src.data_processing.resize_image import resize_images_and_annotations
-from src.data_processing.create_config import create_config_data, delete_cache
-from src.data_processing.convert_yolo_to_voc import convert_voc
-from src.train_utils.train_model.scripts.yolov5_train import train_V5
-from src.train_utils.train_model.scripts.yolov7_train import train_V7
-from src.train_utils.train_model.scripts.yolov8_train import train_V8
-from src.train_utils.train_model.scripts.faster_rccn_train import train_frcnn
-from src.train_utils.train_model.scripts.ssd_train import train_ssd
-from src.data_processing.utils import modelSelection, loadConfig, getDataPath, getClassesPath
+from src.data_processing.data_utils.utils import load_config, get_data_path
+from src.data_processing.data_utils.split_dataset import split_data, copy_arch_folder, resize_images_and_annotations
+from src.data_processing.train_processing.prepare_train import get_classes_path, model_selection, delete_cache
+from src.data_processing.train_processing.prepare_train import create_config_data, check_config_arrays_sizes
+from src.data_processing.train_processing.convert_yolo_to_voc import convert_voc
+
+from src.train_utils.train_models.scripts import yolov8_train, yolov7_train, yolov5_train
+from src.train_utils.train_models.scripts import faster_rccn_train, ssd_train
+
+
 
 
 FILE = Path(__file__).resolve()
@@ -22,48 +22,52 @@
     sys.path.append(str(ROOT))
 
 
-def fit_model(DATA_PATH, CLASSES, IMG_SIZE, BATCH_SIZE, EPOCHS, MODEL, CONFIG_PATH, SPLIT_TRAIN_VALUE,
+def fit_model(DATA_PATH, CLASSES, IMG_SIZE, BATCH_SIZE, EPOCHS, MODEL, CONFIG_NAME, SPLIT_TRAIN_VALUE,
               SPLIT_VAL_VALUE, GPU_COUNT, SELECT_GPU):
 
-    DATA_PATH = getDataPath(ROOT, DATA_PATH)
-    CLASSES = getClassesPath(ROOT, CLASSES)
+    DATA_PATH = get_data_path(ROOT, DATA_PATH)
+    CLASSES_PATH = get_classes_path(ROOT, CLASSES)
 
     PATH_SPLIT_TRAIN, PATH_SPLIT_VALID = split_data(DATA_PATH, SPLIT_TRAIN_VALUE, SPLIT_VAL_VALUE)
 
-    arch, MODEL_PATH = modelSelection(MODEL)
+    ARCH, MODEL_PATH = model_selection(MODEL)
 
     delete_cache(DATA_PATH)
-    #ready
-    if arch == 'yolov8':
-        CONFIG_PATH = create_config_data(PATH_SPLIT_TRAIN, PATH_SPLIT_VALID, CLASSES, CONFIG_PATH, arch, BATCH_SIZE,
+
+    if ARCH == 'yolov8':
+        CONFIG_PATH = create_config_data(PATH_SPLIT_TRAIN, PATH_SPLIT_VALID, CLASSES_PATH, CONFIG_NAME, ARCH, BATCH_SIZE,
                                          EPOCHS, MODEL)
-        train_V8(IMG_SIZE, BATCH_SIZE, EPOCHS, CONFIG_PATH, MODEL_PATH, GPU_COUNT, SELECT_GPU)
-    #ready
-    elif arch == 'yolov5':
-        CONFIG_PATH = create_config_data(PATH_SPLIT_TRAIN, PATH_SPLIT_VALID, CLASSES, CONFIG_PATH, arch, BATCH_SIZE,
+        yolov8_train.train_V8(IMG_SIZE, BATCH_SIZE, EPOCHS, CONFIG_PATH, MODEL_PATH, GPU_COUNT, SELECT_GPU)
+    elif ARCH == 'yolov5':
+        CONFIG_PATH = create_config_data(PATH_SPLIT_TRAIN, PATH_SPLIT_VALID, CLASSES_PATH, CONFIG_NAME, ARCH, BATCH_SIZE,
                                          EPOCHS, MODEL)
-        train_V5(IMG_SIZE, BATCH_SIZE, EPOCHS, CONFIG_PATH, MODEL_PATH, GPU_COUNT, SELECT_GPU)
-    #ready
-    elif arch == 'yolov7':
-        CONFIG_PATH = create_config_data(PATH_SPLIT_TRAIN, PATH_SPLIT_VALID, CLASSES, CONFIG_PATH, arch, BATCH_SIZE,
+        yolov5_train.train_V5(IMG_SIZE, BATCH_SIZE, EPOCHS, CONFIG_PATH, MODEL_PATH, GPU_COUNT, SELECT_GPU)
+    elif ARCH == 'yolov7':
+        CONFIG_PATH = create_config_data(PATH_SPLIT_TRAIN, PATH_SPLIT_VALID, CLASSES_PATH, CONFIG_NAME, ARCH, BATCH_SIZE,
                                          EPOCHS, MODEL)
-        train_V7(IMG_SIZE, BATCH_SIZE, EPOCHS, CONFIG_PATH, MODEL_PATH, GPU_COUNT, SELECT_GPU)
-    #ready
-    elif arch == 'faster-rcnn':
+        yolov7_train.train_V7(IMG_SIZE, BATCH_SIZE, EPOCHS, CONFIG_PATH, MODEL_PATH, GPU_COUNT, SELECT_GPU)
+    elif ARCH == 'faster-rcnn':
+
         DATA_PATH = copy_arch_folder(DATA_PATH)
+        PATH_SPLIT_TRAIN = Path(DATA_PATH) / 'train'
+        PATH_SPLIT_VALID = Path(DATA_PATH) / 'valid'
         resize_images_and_annotations(DATA_PATH, IMG_SIZE)
         convert_voc(DATA_PATH, CLASSES)
-        CONFIG_PATH = create_config_data(Path(DATA_PATH) / 'train', Path(DATA_PATH) / 'valid', CLASSES, CONFIG_PATH, arch,
+
+        CONFIG_PATH = create_config_data(PATH_SPLIT_TRAIN, PATH_SPLIT_VALID, CLASSES_PATH, CONFIG_NAME, ARCH,
                                          BATCH_SIZE, EPOCHS, MODEL)
-        train_frcnn(CONFIG_PATH, EPOCHS, BATCH_SIZE, GPU_COUNT, IMG_SIZE)
-    #ready
-    elif arch == 'ssd':
+        faster_rccn_train.train_frcnn(CONFIG_PATH, EPOCHS, BATCH_SIZE, GPU_COUNT, IMG_SIZE)
+    elif ARCH == 'ssd':
+
         DATA_PATH = copy_arch_folder(DATA_PATH)
+        PATH_SPLIT_TRAIN = Path(DATA_PATH) / 'train.json'
+        PATH_SPLIT_VALID = Path(DATA_PATH) / 'valid.json'
         resize_images_and_annotations(DATA_PATH, IMG_SIZE)
         convert_voc(DATA_PATH, CLASSES)
-        CONFIG_PATH = create_config_data(Path(DATA_PATH) / 'train.json', Path(DATA_PATH) / 'valid.json', CLASSES, CONFIG_PATH,
-                                         arch, BATCH_SIZE, EPOCHS, MODEL)
-        train_ssd(CONFIG_PATH)
+
+        CONFIG_PATH = create_config_data(PATH_SPLIT_TRAIN, PATH_SPLIT_VALID, CLASSES, CONFIG_NAME,
+                                         ARCH, BATCH_SIZE, EPOCHS, MODEL)
+        ssd_train.train_ssd(CONFIG_PATH)
 
 
 def prepare_to_train(config, list_parameters):
@@ -73,7 +77,7 @@ def prepare_to_train(config, list_parameters):
     BATCH_SIZE = config['BATCH_SIZE']
     EPOCHS = config['EPOCHS']
     MODEL = config['MODEL']
-    CONFIG_PATH = config['CONFIG_PATH']
+    CONFIG_NAME = 'dataset.yaml'
     SPLIT_TRAIN_VALUE = config['SPLIT_TRAIN_VALUE']
     SPLIT_VAL_VALUE = config['SPLIT_VAL_VALUE']
     GPU_COUNT = config['GPU_COUNT']
@@ -88,7 +92,7 @@ def prepare_to_train(config, list_parameters):
                 'BATCH_SIZE': BATCH_SIZE[i] if isinstance(BATCH_SIZE, list) else BATCH_SIZE,
                 'EPOCHS': EPOCHS[i] if isinstance(EPOCHS, list) else EPOCHS,
                 'MODEL': MODEL[i] if isinstance(MODEL, list) else MODEL,
-                'CONFIG_PATH': CONFIG_PATH[i] if isinstance(CONFIG_PATH, list) else CONFIG_PATH,
+                'CONFIG_NAME': CONFIG_NAME[i] if isinstance(CONFIG_NAME, list) else CONFIG_NAME,
                 'SPLIT_TRAIN_VALUE': SPLIT_TRAIN_VALUE[i] if isinstance(SPLIT_TRAIN_VALUE, list) else SPLIT_TRAIN_VALUE,
                 'SPLIT_VAL_VALUE': SPLIT_VAL_VALUE[i] if isinstance(SPLIT_VAL_VALUE, list) else SPLIT_VAL_VALUE,
                 'GPU_COUNT': GPU_COUNT[i] if isinstance(GPU_COUNT, list) else GPU_COUNT,
@@ -97,26 +101,16 @@ def prepare_to_train(config, list_parameters):
             fit_model(**current_params)
 
     else:
-        fit_model(DATA_PATH, CLASSES, IMG_SIZE, BATCH_SIZE, EPOCHS, MODEL, CONFIG_PATH, SPLIT_TRAIN_VALUE,
+        fit_model(DATA_PATH, CLASSES, IMG_SIZE, BATCH_SIZE, EPOCHS, MODEL, CONFIG_NAME, SPLIT_TRAIN_VALUE,
                 SPLIT_VAL_VALUE, GPU_COUNT, SELECT_GPU)
 
 
-def check_dict_arrays_sizes(dictionary):
-    for key, value in dictionary.items():
-        if isinstance(value, list):
-            first_array = next(iter(dictionary.values()))
-            first_array_size = len(first_array)
-            current_array_size = len(value)
-            if current_array_size != first_array_size:
-                raise ValueError(f"Size mismatch for key '{key}'. Expected size: {first_array_size}, actual size: {current_array_size}")
-
-
 def run():
-    config_path = Path(ROOT) / 'ODRS' / 'train_utils' / 'config' / 'custom_config.yaml'
-    config = loadConfig(config_path)
+    config_path = Path(ROOT) / 'src' / 'train_utils' / 'config' / 'train_config.yaml'
+    config = load_config(config_path)
 
     list_parameters = {key: value for key, value in config.items() if isinstance(value, list)}
-    check_dict_arrays_sizes(list_parameters)
+    check_config_arrays_sizes(list_parameters)
     prepare_to_train(config, list_parameters)