diff --git a/examples/README.md b/examples/welding-defect-detection/README.md similarity index 100% rename from examples/README.md rename to examples/welding-defect-detection/README.md diff --git a/examples/welding-defect-detection/client/custom.yaml b/examples/welding-defect-detection/client/custom.yaml new file mode 100644 index 000000000..f37044c8e --- /dev/null +++ b/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) diff --git a/examples/welding-defect-detection/client/data.py b/examples/welding-defect-detection/client/data.py new file mode 100644 index 000000000..d2aea5c69 --- /dev/null +++ b/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() diff --git a/examples/welding-defect-detection/client/fedn.yaml b/examples/welding-defect-detection/client/fedn.yaml new file mode 100644 index 000000000..526b4a428 --- /dev/null +++ b/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 + \ No newline at end of file diff --git a/examples/welding-defect-detection/client/model.py b/examples/welding-defect-detection/client/model.py new file mode 100644 index 000000000..3272d06f4 --- /dev/null +++ b/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") diff --git a/examples/welding-defect-detection/client/python_env.yaml b/examples/welding-defect-detection/client/python_env.yaml new file mode 100644 index 000000000..b10d2d1a3 --- /dev/null +++ b/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 diff --git a/examples/welding-defect-detection/client/train.py b/examples/welding-defect-detection/client/train.py new file mode 100644 index 000000000..ed9390735 --- /dev/null +++ b/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]) diff --git a/examples/welding-defect-detection/client/validate.py b/examples/welding-defect-detection/client/validate.py new file mode 100644 index 000000000..50dc8a4d0 --- /dev/null +++ b/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])