henry000 commited on
Commit
b80dc1e
·
1 Parent(s): 97fb10e

🔨 [Init] the base train code, and load yaml cfg

Browse files
Files changed (4) hide show
  1. config/model/v7-base.yaml +239 -0
  2. model/yolo.py +35 -0
  3. train.py +26 -0
  4. utils/tools.py +67 -0
config/model/v7-base.yaml ADDED
@@ -0,0 +1,239 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ anchor:
2
+ [1, 2, 3]
3
+ model:
4
+ backbone:
5
+ - Conv:
6
+ args: {out_channels: 32, kernel_size: 3}
7
+ - Conv:
8
+ args: {out_channels: 64, kernel_size: 3, stride: 2}
9
+ - Conv:
10
+ args: {out_channels: 64, kernel_size: 3}
11
+ - Conv:
12
+ args: {out_channels: 128, kernel_size: 3, stride: 2}
13
+ - Conv:
14
+ args: {out_channels: 64, kernel_size: 1}
15
+ - Conv:
16
+ args: {out_channels: 64, kernel_size: 1}
17
+ source: -2
18
+ - Conv:
19
+ args: {out_channels: 64, kernel_size: 3}
20
+ - Conv:
21
+ args: {out_channels: 64, kernel_size: 3}
22
+ - Conv:
23
+ args: {out_channels: 64, kernel_size: 3}
24
+ - Conv:
25
+ args: {out_channels: 64, kernel_size: 3}
26
+ - Concat:
27
+ source: [-1, -3, -5, -6]
28
+ - Conv:
29
+ args: {out_channels: 256, kernel_size: 1}
30
+ - MP:
31
+ args: []
32
+ - Conv:
33
+ args: {out_channels: 128, kernel_size: 1}
34
+ - Conv:
35
+ args: {out_channels: 128, kernel_size: 1}
36
+ source: -3
37
+ - Conv:
38
+ args: {out_channels: 128, kernel_size: 3, stride: 2}
39
+ - Concat:
40
+ source: [-1, -3]
41
+ - Conv:
42
+ args: {out_channels: 128, kernel_size: 1}
43
+ - Conv:
44
+ args: {out_channels: 128, kernel_size: 1}
45
+ source: -2
46
+ - Conv:
47
+ args: {out_channels: 128, kernel_size: 3}
48
+ - Conv:
49
+ args: {out_channels: 128, kernel_size: 3}
50
+ - Conv:
51
+ args: {out_channels: 128, kernel_size: 3}
52
+ - Conv:
53
+ args: {out_channels: 128, kernel_size: 3}
54
+ - Concat:
55
+ source: [-1, -3, -5, -6]
56
+ tags: 8x
57
+ - Conv:
58
+ args: {out_channels: 512, kernel_size: 1}
59
+ - MP:
60
+ args: []
61
+ - Conv:
62
+ args: {out_channels: 256, kernel_size: 1}
63
+ - Conv:
64
+ args: {out_channels: 256, kernel_size: 1}
65
+ source: -3
66
+ - Conv:
67
+ args: {out_channels: 256, kernel_size: 3, stride: 2}
68
+ - Concat:
69
+ source: [-1, -3]
70
+ - Conv:
71
+ args: {out_channels: 256, kernel_size: 1}
72
+ - Conv:
73
+ args: {out_channels: 256, kernel_size: 1}
74
+ source: -2
75
+ - Conv:
76
+ args: {out_channels: 256, kernel_size: 3}
77
+ - Conv:
78
+ args: {out_channels: 256, kernel_size: 3}
79
+ - Conv:
80
+ args: {out_channels: 256, kernel_size: 3}
81
+ - Conv:
82
+ args: {out_channels: 256, kernel_size: 3}
83
+ - Concat:
84
+ source: [-1, -3, -5, -6]
85
+ - Conv:
86
+ args: {out_channels: 1024, kernel_size: 1}
87
+ tags: 16x
88
+ - MP:
89
+ args: []
90
+ - Conv:
91
+ args: {out_channels: 512, kernel_size: 1}
92
+ - Conv:
93
+ args: {out_channels: 512, kernel_size: 1}
94
+ source: -3
95
+ - Conv:
96
+ args: {out_channels: 512, kernel_size: 3, stride: 2}
97
+ - Concat:
98
+ source: [-1, -3]
99
+ - Conv:
100
+ args: {out_channels: 256, kernel_size: 1}
101
+ - Conv:
102
+ args: {out_channels: 256, kernel_size: 1}
103
+ source: -2
104
+ - Conv:
105
+ args: {out_channels: 256, kernel_size: 3}
106
+ - Conv:
107
+ args: {out_channels: 256, kernel_size: 3}
108
+ - Conv:
109
+ args: {out_channels: 256, kernel_size: 3}
110
+ - Conv:
111
+ args: {out_channels: 256, kernel_size: 3}
112
+ - Concat:
113
+ source: [-1, -3, -5, -6]
114
+ - Conv:
115
+ args: {out_channels: 1024, kernel_size: 1}
116
+ tags: 32x
117
+ head:
118
+ - SPPCSPC:
119
+ args: [512]
120
+ - Conv:
121
+ args: {out_channels: 256, kernel_size: 1}
122
+ - Upsample:
123
+ args: [None, 2, nearest]
124
+ - Conv:
125
+ args: {out_channels: 256, kernel_size: 1}
126
+ source: 16x
127
+ - Concat:
128
+ source: [-1, -2]
129
+ - Conv:
130
+ args: {out_channels: 256, kernel_size: 1}
131
+ - Conv:
132
+ args: {out_channels: 256, kernel_size: 1}
133
+ source: -2
134
+ - Conv:
135
+ args: {out_channels: 128, kernel_size: 3}
136
+ - Conv:
137
+ args: {out_channels: 128, kernel_size: 3}
138
+ - Conv:
139
+ args: {out_channels: 128, kernel_size: 3}
140
+ - Conv:
141
+ args: {out_channels: 128, kernel_size: 3}
142
+ - Concat:
143
+ source: [-1, -2, -3, -4, -5, -6]
144
+ - Conv:
145
+ args: {out_channels: 256, kernel_size: 1}
146
+ - Conv:
147
+ args: {out_channels: 128, kernel_size: 1}
148
+ - Upsample:
149
+ args: [None, 2, nearest]
150
+ - Conv:
151
+ args: {out_channels: 128, kernel_size: 1}
152
+ source: 8x
153
+ - Concat:
154
+ source: [-1, -2]
155
+ - Conv:
156
+ args: {out_channels: 128, kernel_size: 1}
157
+ - Conv:
158
+ args: {out_channels: 128, kernel_size: 1}
159
+ source: -2
160
+ - Conv:
161
+ args: {out_channels: 64, kernel_size: 3}
162
+ - Conv:
163
+ args: {out_channels: 64, kernel_size: 3}
164
+ - Conv:
165
+ args: {out_channels: 64, kernel_size: 3}
166
+ - Conv:
167
+ args: {out_channels: 64, kernel_size: 3}
168
+ - Concat:
169
+ source: [-1, -2, -3, -4, -5, -6]
170
+ - Conv:
171
+ args: {out_channels: 128, kernel_size: 1}
172
+ - MP:
173
+ args: []
174
+ - Conv:
175
+ args: {out_channels: 128, kernel_size: 1}
176
+ - Conv:
177
+ args: {out_channels: 128, kernel_size: 1}
178
+ source: -3
179
+ - Conv:
180
+ args: {out_channels: 128, kernel_size: 3, stride: 2}
181
+ - Concat:
182
+ source: [-1, -3, 63]
183
+ - Conv:
184
+ args: {out_channels: 256, kernel_size: 1}
185
+ - Conv:
186
+ args: {out_channels: 256, kernel_size: 1}
187
+ source: -2
188
+ - Conv:
189
+ args: {out_channels: 128, kernel_size: 3}
190
+ - Conv:
191
+ args: {out_channels: 128, kernel_size: 3}
192
+ - Conv:
193
+ args: {out_channels: 128, kernel_size: 3}
194
+ - Conv:
195
+ args: {out_channels: 128, kernel_size: 3}
196
+ - Concat:
197
+ source: [-1, -2, -3, -4, -5, -6]
198
+ - Conv:
199
+ args: {out_channels: 256, kernel_size: 1}
200
+ - MP:
201
+ args: []
202
+ - Conv:
203
+ args: {out_channels: 256, kernel_size: 1}
204
+ - Conv:
205
+ args: {out_channels: 256, kernel_size: 1}
206
+ source: -3
207
+ - Conv:
208
+ args: {out_channels: 256, kernel_size: 3, stride: 2}
209
+ - Concat:
210
+ source: [-1, -3, 51]
211
+ - Conv:
212
+ args: {out_channels: 512, kernel_size: 1}
213
+ - Conv:
214
+ args: {out_channels: 512, kernel_size: 1}
215
+ source: -2
216
+ - Conv:
217
+ args: {out_channels: 256, kernel_size: 3}
218
+ - Conv:
219
+ args: {out_channels: 256, kernel_size: 3}
220
+ - Conv:
221
+ args: {out_channels: 256, kernel_size: 3}
222
+ - Conv:
223
+ args: {out_channels: 256, kernel_size: 3}
224
+ - Concat:
225
+ source: [-1, -2, -3, -4, -5, -6]
226
+ - Conv:
227
+ args: {out_channels: 512, kernel_size: 1}
228
+ - RepConv:
229
+ args: [256, 3, 1]
230
+ source: 75
231
+ - RepConv:
232
+ args: [512, 3, 1]
233
+ source: 88
234
+ - RepConv:
235
+ args: [1024, 3, 1]
236
+ source: 101
237
+ - IDetect:
238
+ args: [nc, anchors]
239
+ source: [102, 103, 104]
model/yolo.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch.nn as nn
2
+ from loguru import logger
3
+ from typing import Dict, Any
4
+
5
+
6
+ class YOLO(nn.Module):
7
+ """
8
+ A preliminary YOLO (You Only Look Once) model class still under development.
9
+
10
+ This class is intended to define a YOLO model for object detection tasks. It is
11
+ currently not implemented and serves as a placeholder for future development.
12
+
13
+ Parameters:
14
+ model_cfg: Configuration for the YOLO model. Expected to define the layers,
15
+ parameters, and any other relevant configuration details.
16
+ """
17
+
18
+ def __init__(self, model_cfg: Dict[str, Any]):
19
+ super(YOLO, self).__init__()
20
+ # Placeholder for initialization logic
21
+ print(model_cfg)
22
+ raise NotImplementedError("Constructor not implemented.")
23
+
24
+
25
+ def get_model(model_cfg: dict) -> YOLO:
26
+ """Constructs and returns a model from a Dictionary configuration file.
27
+
28
+ Args:
29
+ config_file (dict): The configuration file of the model.
30
+
31
+ Returns:
32
+ YOLO: An instance of the model defined by the given configuration.
33
+ """
34
+ model = YOLO(model_cfg)
35
+ return model
train.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import sys
3
+ from loguru import logger
4
+ from model.yolo import get_model
5
+ from utils.tools import load_model_cfg
6
+
7
+
8
+ def parse_arguments() -> argparse.Namespace:
9
+ """
10
+ Parse command-line arguments to get the model configuration file.
11
+
12
+ Returns:
13
+ argparse.Namespace: The command-line arguments object with 'config' attribute.
14
+ """
15
+ parser = argparse.ArgumentParser(description="Load a YOLO model configuration and display the model.")
16
+ parser.add_argument(
17
+ "--model-config", type=str, default="v7-base.yaml", help="Name or path to the model configuration file."
18
+ )
19
+ return parser.parse_args()
20
+
21
+
22
+ if __name__ == "__main__":
23
+ args = parse_arguments()
24
+ model_cfg = load_model_cfg(args.model_config)
25
+ model = get_model(model_cfg)
26
+ logger.info("Success load model: {}", model)
utils/tools.py ADDED
@@ -0,0 +1,67 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import yaml
3
+ from loguru import logger
4
+ from typing import Dict, Any
5
+
6
+
7
+ def complete_path(file_name: str = "v7-base.yaml") -> str:
8
+ """
9
+ Ensures the path to a model configuration is a existing file
10
+
11
+ Parameters:
12
+ file_name (str): The filename or path, with default 'v7-base.yaml'.
13
+
14
+ Returns:
15
+ str: A complete path with necessary prefix and extension.
16
+ """
17
+ # Ensure the file has the '.yaml' extension if missing
18
+ if not file_name.endswith(".yaml"):
19
+ file_name += ".yaml"
20
+
21
+ # Add folder prefix if only the filename is provided
22
+ if os.path.dirname(file_name) == "":
23
+ file_name = os.path.join("./config/model", file_name)
24
+
25
+ return file_name
26
+
27
+
28
+ def load_model_cfg(file_path: str) -> Dict[str, Any]:
29
+ """
30
+ Read a YAML configuration file, ensure necessary keys are present, and return its content as a dictionary.
31
+
32
+ Args:
33
+ file_path (str): The path to the YAML configuration file.
34
+
35
+ Returns:
36
+ Dict[str, Any]: The contents of the YAML file as a dictionary.
37
+
38
+ Raises:
39
+ FileNotFoundError: If the YAML file cannot be found.
40
+ yaml.YAMLError: If there is an error parsing the YAML file.
41
+ """
42
+ file_path = complete_path(file_path)
43
+ try:
44
+ with open(file_path, "r") as file:
45
+ model_cfg = yaml.safe_load(file) or {}
46
+
47
+ # Check for required keys and set defaults if not present
48
+ if "nc" not in model_cfg:
49
+ model_cfg["nc"] = 80
50
+ logger.warning("'nc' not found in the YAML file. Setting default 'nc' to 80.")
51
+
52
+ if "anchor" not in model_cfg:
53
+ logger.error("'anchor' is missing in the configuration file.")
54
+ raise ValueError("Missing required key: 'anchor'")
55
+
56
+ if "model" not in model_cfg:
57
+ logger.error("'model' is missing in the configuration file.")
58
+ raise ValueError("Missing required key: 'model'")
59
+
60
+ return model_cfg
61
+
62
+ except FileNotFoundError:
63
+ logger.error(f"YAML file not found: {file_path}")
64
+ raise
65
+ except yaml.YAMLError as e:
66
+ logger.error(f"Error parsing YAML file: {e}")
67
+ raise