Feature/SK-971 | New object detection example

examples/welding-defect-detection/client/custom.yaml

@@ -0,0 +1,46 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
+
+# Parameters
+nc: 3 # number of classes
+scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
+  # [depth, width, max_channels]
+  n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers,  3157200 parameters,  3157184 gradients,   8.9 GFLOPs
+  s: [0.33, 0.50, 1024] # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients,  28.8 GFLOPs
+  m: [0.67, 0.75, 768] # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients,  79.3 GFLOPs
+  l: [1.00, 1.00, 512] # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs
+  x: [1.00, 1.25, 512] # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs
+
+# YOLOv8.0n backbone
+backbone:
+  # [from, repeats, module, args]
+  - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
+  - [-1, 3, C2f, [128, True]]
+  - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
+  - [-1, 6, C2f, [256, True]]
+  - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
+  - [-1, 6, C2f, [512, True]]
+  - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
+  - [-1, 3, C2f, [1024, True]]
+  - [-1, 1, SPPF, [1024, 5]] # 9
+
+# YOLOv8.0n head
+head:
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
+  - [[-1, 6], 1, Concat, [1]] # cat backbone P4
+  - [-1, 3, C2f, [512]] # 12
+
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
+  - [[-1, 4], 1, Concat, [1]] # cat backbone P3
+  - [-1, 3, C2f, [256]] # 15 (P3/8-small)
+
+  - [-1, 1, Conv, [256, 3, 2]]
+  - [[-1, 12], 1, Concat, [1]] # cat head P4
+  - [-1, 3, C2f, [512]] # 18 (P4/16-medium)
+
+  - [-1, 1, Conv, [512, 3, 2]]
+  - [[-1, 9], 1, Concat, [1]] # cat head P5
+  - [-1, 3, C2f, [1024]] # 21 (P5/32-large)
+
+  - [[15, 18, 21], 1, Detect, [nc]] # Detect(P3, P4, P5)

examples/welding-defect-detection/client/data.py

@@ -0,0 +1,99 @@
+import os
+from math import floor
+import torch
+import yaml
+import opendatasets
+#from sklearn import preprocessing
+dir_path = os.path.dirname(os.path.realpath(__file__))
+abs_path = os.path.abspath(dir_path)
+
+
+def get_data(out_dir=None):
+
+    # Only download if not already downloaded
+    if not os.path.exists(f"{out_dir}/welding-defect-object-detection"):
+        opendatasets.download('https://www.kaggle.com/datasets/sukmaadhiwijaya/welding-defect-object-detection')
+
+
+def load_labels(label_dir):
+    label_files = os.listdir(label_dir)
+    data = []
+    for label_file in label_files:
+        with open(os.path.join(label_dir, label_file), 'r') as file:
+            lines = file.readlines()
+            for line in lines:
+                class_id, x_center, y_center, width, height = map(float, line.strip().split())
+                data.append([class_id, x_center, y_center, width, height])
+    return data
+
+
+def load_data(data_path=None, is_train=True, as_yaml=True):
+    if data_path is None:
+        data_path = os.environ.get("FEDN_DATA_PATH", abs_path + "welding-defect-object-detection/The Welding Defect Dataset/The Welding Defect Dataset")
+
+    yaml = data_path + '/data.yaml'
+    path = None
+    if is_train:
+        path = data_path + "/train/images"
+    else:
+        path = data_path + "/test/images"    
+    dir_list = os.listdir(path)
+
+    if as_yaml:
+        return yaml, len(dir_list)
+    else:
+        return dir_list, len(dir_list)
+
+
+def splitset(dataset, parts):
+
+    n = dataset.shape[0]
+    local_n = floor(n / parts)
+    result = []
+    for i in range(parts):
+        result.append(dataset[i * local_n : (i + 1) * local_n])
+    return result
+
+
+
+def split(out_dir="package"):
+    n_splits = int(os.environ.get("FEDN_NUM_DATA_SPLITS", 1))
+
+    # Make dir
+    if not os.path.exists(f"{out_dir}/client"):
+        os.mkdir(f"{out_dir}/client")
+
+    # Load and convert to dict
+    X_train = load_data(is_train=True, as_yaml=False)
+    X_test = load_data(is_train=False, as_yaml=False)
+
+    y_train = load_labels(abs_path + "welding-defect-object-detection/The Welding Defect Dataset/The Welding Defect Dataset/train/labels")
+    y_test = load_labels(abs_path + "welding-defect-object-detection/The Welding Defect Dataset/The Welding Defect Dataset/test/labels")
+
+    data = {
+        "x_train": splitset(X_train, n_splits),
+        "y_train": splitset(y_train, n_splits),
+        "x_test": splitset(X_test, n_splits),
+        "y_test": splitset(y_test, n_splits),
+    }
+
+    # Make splits
+    for i in range(n_splits):
+        subdir = f"{out_dir}/client/{str(i+1)}"
+        if not os.path.exists(subdir):
+            os.mkdir(subdir)
+        torch.save(
+            {
+                "x_train": data["x_train"][i],
+                "y_train": data["y_train"][i],
+                "x_test": data["x_test"][i],
+                "y_test": data["y_test"][i],
+            },
+            f"{subdir}/welding.pt",
+        )
+
+
+if __name__ == "__main__":
+
+    get_data()
+    split()

examples/welding-defect-detection/client/fedn.yaml

@@ -0,0 +1,11 @@
+python_env: python_env.yaml
+entry_points:
+  build:
+    command: python model.py
+  startup:
+    command: python data.py
+  train:
+    command: python train.py
+  validate:
+    command: python validate.py
+

examples/welding-defect-detection/client/model.py

@@ -0,0 +1,71 @@
+import collections
+from ultralytics import YOLO
+import torch
+
+from fedn.utils.helpers.helpers import get_helper
+
+HELPER_MODULE = "numpyhelper"
+helper = get_helper(HELPER_MODULE)
+
+
+def compile_model():
+    """Compile the pytorch model.
+
+    :return: The compiled model.
+    :rtype: torch.nn.Module
+    """
+    model = YOLO('custom.yaml')
+
+    return model
+
+
+def save_parameters(model, out_path):
+    """Save model paramters to file.
+
+    :param model: The model to serialize.
+    :type model: torch.nn.Module
+    :param out_path: The path to save to.
+    :type out_path: str
+    """
+    parameters_np = [val.cpu().numpy() for _, val in model.state_dict().items()]
+
+    helper.save(parameters_np, out_path)
+
+
+def load_parameters(model_path):
+    """Load model parameters from file and populate model.
+
+    param model_path: The path to load from.
+    :type model_path: str
+    :return: The loaded model.
+    :rtype: torch.nn.Module
+    """
+    model = compile_model()
+    parameters_np = helper.load(model_path)
+
+    params_dict = zip(model.state_dict().keys(), parameters_np)
+    state_dict = collections.OrderedDict({key: torch.tensor(x) for key, x in params_dict})
+
+    model.load_state_dict(state_dict, strict=True) 
+
+    torch.save(model,'tempfile.pt')
+    model = YOLO('tempfile.pt')
+
+    return model
+
+
+def init_seed(out_path="seed.npz"):
+    """Initialize seed model and save it to file.
+
+
+    :param out_path: The path to save the seed model to.
+    :type out_path: str
+    """
+    # Init and save
+    model = compile_model()
+
+    save_parameters(model, out_path)
+
+
+if __name__ == "__main__":
+    init_seed("../seed.npz")

examples/welding-defect-detection/client/python_env.yaml

@@ -0,0 +1,10 @@
+name: yolov9-imagerecognition
+build_dependencies:
+  - pip
+  - setuptools
+  - wheel
+dependencies:
+  - torch==2.3.1
+  - torchvision==0.18.1
+  - ultralytics
+  - fedn

examples/welding-defect-detection/client/train.py

@@ -0,0 +1,64 @@
+import sys
+from ultralytics import YOLO
+from model import load_parameters, save_parameters
+from data import load_data
+from fedn.utils.helpers.helpers import save_metadata
+import os
+
+# Get the list of all files and directories
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(os.path.abspath(dir_path))
+
+
+def train(in_model_path, out_model_path, data_path=None, batch_size=64, epochs=1, lr=0.01):
+    """Complete a model update.
+
+    Load model paramters from in_model_path (managed by the FEDn client),
+    perform a model update, and write updated paramters
+    to out_model_path (picked up by the FEDn client).
+
+    :param in_model_path: The path to the input model.
+    :type in_model_path: str
+    :param out_model_path: The path to save the output model to.
+    :type out_model_path: str
+    :param data_path: The path to the data file.
+    :type data_path: str
+    :param batch_size: The batch size to use.
+    :type batch_size: int
+    :param epochs: The number of epochs to train.
+    :type epochs: int
+    :param lr: The learning rate to use.
+    :type lr: float
+    """
+    # Load data
+    data, data_len = load_data(None, is_train=True)
+
+    # Load parmeters and initialize model
+    model = load_parameters(in_model_path)
+
+
+    # Train
+    model.train(data=data, epochs=epochs, imgsz=640, batch=batch_size,
+               lr0=lr, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1,
+               box=0.05, cls=0.5, iou=0.2, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, translate=0.1, scale=0.5, mosaic=1.0, mixup=0.5)
+
+
+    # Metadata needed for aggregation server side
+    metadata = {
+        # num_examples are mandatory
+        "num_examples": data_len,
+        "batch_size": batch_size,
+        "epochs": epochs,
+        "lr": lr,
+    }
+
+    # Save JSON metadata file (mandatory)
+    save_metadata(metadata, out_model_path)
+
+    # Save model update (mandatory)
+    save_parameters(model, out_model_path)
+
+
+if __name__ == "__main__":
+    train(sys.argv[1], sys.argv[2])

examples/welding-defect-detection/client/validate.py

@@ -0,0 +1,47 @@
+import os
+import sys
+
+from model import load_parameters
+
+from data import load_data, load_labels
+from fedn.utils.helpers.helpers import save_metrics
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(os.path.abspath(dir_path))
+
+
+def validate(in_model_path, out_json_path, data_path=None):
+    """Validate model.
+
+    :param in_model_path: The path to the input model.
+    :type in_model_path: str
+    :param out_json_path: The path to save the output JSON to.
+    :type out_json_path: str
+    :param data_path: The path to the data file.
+    :type data_path: str
+    """
+    # Load data
+    train_data_yaml, train_data_length = load_data(None, is_train=True)
+    test_data_yaml, test_data_length = load_data(None, is_train=False)
+
+
+    model = load_parameters(in_model_path)
+
+    validation_results = model.val(data=test_data_yaml)
+
+
+    # JSON schema
+    report = {
+        "map50-95": float(validation_results.box.map),  # map50-95
+        "map50": float(validation_results.box.map50),  # map50
+        "map75": float(validation_results.box.map75),  # map75
+    }
+
+    # Save JSON
+    save_metrics(report, out_json_path)    
+
+
+
+
+if __name__ == "__main__":
+    validate(sys.argv[1], sys.argv[2])