Skip to content

Commit

Permalink
add example
Browse files Browse the repository at this point in the history
  • Loading branch information
KatHellg committed Sep 17, 2024
1 parent 3619bb0 commit 0026e82
Show file tree
Hide file tree
Showing 8 changed files with 348 additions and 0 deletions.
File renamed without changes.
46 changes: 46 additions & 0 deletions examples/welding-defect-detection/client/custom.yaml
Original file line number Diff line number Diff line change
@@ -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)
99 changes: 99 additions & 0 deletions examples/welding-defect-detection/client/data.py
Original file line number Diff line number Diff line change
@@ -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()
11 changes: 11 additions & 0 deletions examples/welding-defect-detection/client/fedn.yaml
Original file line number Diff line number Diff line change
@@ -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

71 changes: 71 additions & 0 deletions examples/welding-defect-detection/client/model.py
Original file line number Diff line number Diff line change
@@ -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")
10 changes: 10 additions & 0 deletions examples/welding-defect-detection/client/python_env.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
name: yolov9-imagerecognition
build_dependencies:
- pip
- setuptools
- wheel
dependencies:
- torch==2.3.1
- torchvision==0.18.1
- ultralytics
- fedn
64 changes: 64 additions & 0 deletions examples/welding-defect-detection/client/train.py
Original file line number Diff line number Diff line change
@@ -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])
47 changes: 47 additions & 0 deletions examples/welding-defect-detection/client/validate.py
Original file line number Diff line number Diff line change
@@ -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])

0 comments on commit 0026e82

Please sign in to comment.