1+ import os
2+ import time
3+ import torch
4+ from mltu .preprocessors import ImageReader
5+ from mltu .annotations .images import CVImage
6+ from mltu .transformers import ImageResizer , ImageShowCV2 , ImageNormalizer
7+ from mltu .augmentors import RandomBrightness , RandomRotate , RandomErodeDilate , RandomSharpen , \
8+ RandomMirror , RandomFlip , RandomGaussianBlur , RandomSaltAndPepper , RandomDropBlock , RandomMosaic
9+ from mltu .torch .model import Model
10+ from mltu .torch .dataProvider import DataProvider
11+ from mltu .torch .yolo .annotation import VOCAnnotationReader
12+ from mltu .torch .yolo .preprocessors import YoloPreprocessor
13+ from mltu .torch .yolo .loss import v8DetectionLoss
14+ from mltu .torch .yolo .metrics import YoloMetrics
15+ from mltu .torch .yolo .optimizer import build_optimizer , AccumulativeOptimizer
16+ from mltu .torch .callbacks import EarlyStopping , ModelCheckpoint , ReduceLROnPlateau , TensorBoard , Model2onnx , WarmupCosineDecay
17+
18+ from ultralytics .nn .tasks import DetectionModel
19+ from ultralytics .engine .model import Model as BaseModel
20+
21+
22+ annotations_path = "Datasets/car-plate-detection/annotations"
23+
24+ dataset = [[None , os .path .join (annotations_path , f )] for f in os .listdir (annotations_path )]
25+
26+ # Make sure torch can see GPU device, it is not recommended to train with CPU
27+ device = torch .device ("cuda" if torch .cuda .is_available () else "cpu" )
28+
29+ img_size = 416
30+ labels = {0 : "licence" }
31+
32+ # Create a data provider for the dataset
33+ data_provider = DataProvider (
34+ dataset = dataset ,
35+ skip_validation = True ,
36+ batch_size = 16 ,
37+ data_preprocessors = [
38+ VOCAnnotationReader (labels = labels ),
39+ ImageReader (CVImage ),
40+ ],
41+ transformers = [
42+ # ImageShowCV2(),
43+ ImageResizer (img_size , img_size ),
44+ ImageNormalizer (transpose_axis = True ),
45+ ],
46+ batch_postprocessors = [
47+ YoloPreprocessor (device , img_size )
48+ ],
49+ numpy = False ,
50+ )
51+
52+ # for b in data_provider:
53+ # pass
54+
55+ # split the dataset into train and test
56+ train_data_provider , val_data_provider = data_provider .split (0.9 , shuffle = False )
57+
58+ # Attaach augmentation to the train data provider
59+ train_data_provider .augmentors = [
60+ RandomBrightness (),
61+ RandomErodeDilate (),
62+ RandomSharpen (),
63+ RandomMirror (),
64+ RandomFlip (),
65+ RandomGaussianBlur (),
66+ RandomSaltAndPepper (),
67+ RandomRotate (angle = 10 ),
68+ RandomDropBlock (),
69+ RandomMosaic (),
70+ ]
71+
72+ # for batch in train_data_provider:
73+ # pass
74+ # print(batch)
75+ # break
76+
77+
78+
79+ base_model = BaseModel ("yolov8n.pt" )
80+ # Create a YOLO model
81+ model = DetectionModel ('yolov8n.yaml' , nc = len (labels ))
82+
83+ try : model .load_state_dict (base_model .model .state_dict (), strict = False )
84+ except : pass
85+
86+ model .to (device )
87+
88+ for k , v in model .named_parameters ():
89+ if any (x in k for x in [".dfl" ]):
90+ print ("freezing" , k )
91+ v .requires_grad = False
92+ elif not v .requires_grad :
93+ v .requires_grad = True
94+
95+ lr = 1e-3
96+ optimizer = build_optimizer (model .model , name = "AdamW" , lr = lr , weight_decay = 0.0 , momentum = 0.937 , decay = 0.0005 )
97+ optimizer = AccumulativeOptimizer (optimizer , 16 , 64 )
98+
99+ # create model object that will handle training and testing of the network
100+ model = Model (
101+ model ,
102+ optimizer ,
103+ v8DetectionLoss (model ),
104+ metrics = [YoloMetrics (nc = len (labels ))],
105+ log_errors = False ,
106+ output_path = f"Models/detector/{ int (time .time ())} ,
107+ clip_grad_norm = 10.0 ,
108+ ema = True ,
109+ )
110+
111+ modelCheckpoint = ModelCheckpoint (monitor = "val_fitness" , mode = "max" , save_best_only = True , verbose = True )
112+ tensorBoard = TensorBoard ()
113+ earlyStopping = EarlyStopping (monitor = "val_fitness" , mode = "max" , patience = 31 , verbose = True )
114+ model2onnx = Model2onnx (input_shape = (1 , 3 , img_size , img_size ), verbose = True , opset_version = 14 ,
115+ dynamic_axes = {"input" : {0 : "batch_size" , 2 : "height" , 3 : "width" },
116+ "output" : {0 : "batch_size" , 2 : "anchors" }},
117+ metadata = {"classes" : labels })
118+ warmupCosineDecayBias = WarmupCosineDecay (lr_after_warmup = lr , final_lr = lr , initial_lr = 0.1 ,
119+ warmup_steps = len (train_data_provider ), warmup_epochs = 10 , ignore_param_groups = [1 , 2 ]) # lr0
120+ warmupCosineDecay = WarmupCosineDecay (lr_after_warmup = lr , final_lr = lr / 10 , initial_lr = 1e-7 ,
121+ warmup_steps = len (train_data_provider ), warmup_epochs = 10 , decay_epochs = 190 , ignore_param_groups = [0 ]) # lr1 and lr2
122+
123+ # Train the model
124+ history = model .fit (
125+ train_data_provider ,
126+ test_dataProvider = val_data_provider ,
127+ epochs = 200 ,
128+ callbacks = [
129+ modelCheckpoint ,
130+ tensorBoard ,
131+ earlyStopping ,
132+ model2onnx ,
133+ warmupCosineDecayBias ,
134+ warmupCosineDecay
135+ ]
136+ )
0 commit comments