File size: 4,236 Bytes
183312f 542860e 183312f b80dc1e 7c6ce21 183312f 97681c2 183312f 97681c2 b80dc1e 97681c2 542860e 7c6ce21 97681c2 7c6ce21 542860e 7c6ce21 97681c2 542860e 97681c2 542860e 97681c2 542860e 7c6ce21 542860e 97681c2 542860e 97681c2 542860e 97681c2 542860e 97681c2 b80dc1e 542860e b80dc1e 97681c2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
import inspect
from typing import Any, Dict, List, Union
import torch
import torch.nn as nn
from loguru import logger
from omegaconf import OmegaConf
from model import module
from utils.tools import load_model_cfg
def get_layer_map():
"""
Dynamically generates a dictionary mapping class names to classes,
filtering to include only those that are subclasses of nn.Module,
ensuring they are relevant neural network layers.
"""
layer_map = {}
for name, obj in inspect.getmembers(module, inspect.isclass):
if issubclass(obj, nn.Module) and obj is not nn.Module:
layer_map[name] = obj
return layer_map
class YOLO(nn.Module):
"""
A preliminary YOLO (You Only Look Once) model class still under development.
Parameters:
model_cfg: Configuration for the YOLO model. Expected to define the layers,
parameters, and any other relevant configuration details.
"""
def __init__(self, model_cfg: Dict[str, Any]):
super(YOLO, self).__init__()
self.nc = model_cfg["nc"]
self.layer_map = get_layer_map() # Get the map Dict[str: Module]
self.build_model(model_cfg.model)
def build_model(self, model_arch: Dict[str, List[Dict[str, Dict[str, Dict]]]]):
model_list = nn.ModuleList()
output_dim = [3]
layer_indices_by_tag = {}
for arch_name in model_arch:
logger.info(f"🏗️ Building model-{arch_name}")
for layer_idx, layer_spec in enumerate(model_arch[arch_name], start=1):
layer_type, layer_info = next(iter(layer_spec.items()))
layer_args = layer_info.get("args", {})
source = layer_info.get("source", -1)
if isinstance(source, str):
source = layer_indices_by_tag[source]
if "Conv" in layer_type:
layer_args["in_channels"] = output_dim[source]
if "Detect" in layer_type:
layer_args["nc"] = self.nc
layer_args["ch"] = [output_dim[idx] for idx in source]
layer = self.create_layer(layer_type, source, **layer_args)
model_list.append(layer)
if "tags" in layer_info:
if layer_info["tags"] in layer_indices_by_tag:
raise ValueError(f"Duplicate tag '{layer_info['tags']}' found.")
layer_indices_by_tag[layer_info["tags"]] = layer_idx
out_channels = self.get_out_channels(layer_type, layer_args, output_dim, source)
output_dim.append(out_channels)
self.model = model_list
def forward(self, x):
y = [x]
for layer in self.model:
if OmegaConf.is_list(layer.source):
model_input = [y[idx] for idx in layer.source]
else:
model_input = y[layer.source]
x = layer(model_input)
y.append(x)
return x
def get_out_channels(self, layer_type: str, layer_args: dict, output_dim: list, source: Union[int, list]):
if "Conv" in layer_type:
return layer_args["out_channels"]
if layer_type in ["MaxPool", "UpSample"]:
return output_dim[source]
if layer_type == "Concat":
return sum(output_dim[idx] for idx in source)
if layer_type == "IDetect":
return None
def create_layer(self, layer_type: str, source: Union[int, list], **kwargs):
if layer_type in self.layer_map:
layer = self.layer_map[layer_type](**kwargs)
layer.source = source
return layer
else:
raise ValueError(f"Unsupported layer type: {layer_type}")
def get_model(model_cfg: dict) -> YOLO:
"""Constructs and returns a model from a Dictionary configuration file.
Args:
config_file (dict): The configuration file of the model.
Returns:
YOLO: An instance of the model defined by the given configuration.
"""
model = YOLO(model_cfg)
logger.info("✅ Success load model")
return model
if __name__ == "__main__":
model_cfg = load_model_cfg("v7-base")
YOLO(model_cfg)
|