🔨 [Init] the base train code, and load yaml cfg
Browse files- config/model/v7-base.yaml +239 -0
- model/yolo.py +35 -0
- train.py +26 -0
- 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
|