-
Notifications
You must be signed in to change notification settings - Fork 36
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
KatHellg
committed
Sep 17, 2024
1 parent
3619bb0
commit 0026e82
Showing
8 changed files
with
348 additions
and
0 deletions.
There are no files selected for viewing
File renamed without changes.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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") |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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]) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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]) |