def __init__(self,
classes=1000,
load_backbone_weights=False,
config_file=None,
weights_file=None):
"""
load the model and the sequential head so that the backbone can be applied for classification
Model tested on ImageNet Tiny
Args:
classes: integer for how many classes can be predicted
load_backbone_weights: bool, if true we will auto load the original darnet weights to use in the model
config_file: str path for the location of the configuration file to use when decoding darknet weights
weights_file: str path with the file containing the dark net weights
"""
super(DarkNet53, self).__init__()
self.backbone = Backbone_Builder("darknet53")
self.head = ks.Sequential([
ks.layers.GlobalAveragePooling2D(),
ks.layers.Dense(classes, activation="sigmoid")
])
if load_backbone_weights:
if config_file is None:
config_file = download('yolov3.cfg')
if weights_file is None:
weights_file = download('yolov3.weights')
full_model = DarkNetConverter.read(config_file, weights_file)
encoder, decoder = split_converter(full_model, 76)
load_weights_dnBackbone(self.backbone, encoder)
return
def load_weights_from_dn(self,
dn2tf_backbone=True,
dn2tf_head=True,
config_file=None,
weights_file=None):
"""
load the entire Yolov3 Model for tensorflow
example:
load yolo with darknet wieghts for backbone
model = Yolov3()
model.build(input_shape = (1, 416, 416, 3))
model.load_weights_from_dn(dn2tf_backbone = True, dn2tf_head = True)
to be implemented
example:
load custom back bone weigths
example:
load custom head weigths
example:
load back bone weigths from tensorflow (our training)
example:
load head weigths from tensorflow (our training)
Args:
dn2tf_backbone: bool, if true it will load backbone weights for yolo v3 from darknet .weights file
dn2tf_head: bool, if true it will load head weights for yolo v3 from darknet .weights file
config_file: str path for the location of the configuration file to use when decoding darknet weights
weights_file: str path with the file containing the dark net weights
"""
if not self._built:
self.build(self._input_shape)
if dn2tf_backbone or dn2tf_neck or dn2tf_head:
if config_file is None:
config_file = download(self._model_name + '.cfg')
if weights_file is None:
weights_file = download(self._model_name + '.weights')
list_encdec = DarkNetConverter.read(config_file, weights_file)
encoder, neck, decoder = split_converter(
list_encdec, self._encoder_decoder_split_location, 138)
if dn2tf_backbone:
#load_weights_dnBackbone(self._backbone, encoder, mtype = self._backbone_name)
load_weights_backbone(self._backbone, encoder)
self._backbone.trainable = False
if dn2tf_head:
load_weights_backbone(self._neck, neck)
self._neck.trainable = False
load_weights_v4head(self._head, decoder)
self._head.trainable = False
return
def initialize(self, model: tf.keras.Model):
if self.task_config.load_darknet_weights:
from yolo.utils import DarkNetConverter
from yolo.utils._darknet2tf.load_weights import split_converter
from yolo.utils._darknet2tf.load_weights2 import load_weights_backbone
from yolo.utils._darknet2tf.load_weights2 import load_head
from yolo.utils._darknet2tf.load_weights2 import load_weights_prediction_layers
from yolo.utils.downloads.file_manager import download
weights_file = self.task_config.model.darknet_weights_file
config_file = self.task_config.model.darknet_weights_cfg
if ('cache' not in weights_file and 'cache' not in config_file):
list_encdec = DarkNetConverter.read(config_file, weights_file)
else:
import os
path = os.path.abspath('cache')
if (not os.path.isdir(path)):
os.mkdir(path)
cfg = f"{path}/cfg/{config_file.split('/')[-1]}"
if not os.path.isfile(cfg):
download(config_file.split('/')[-1])
wgt = f"{path}/weights/{weights_file.split('/')[-1]}"
if not os.path.isfile(wgt):
download(weights_file.split('/')[-1])
list_encdec = DarkNetConverter.read(cfg, wgt)
splits = model.backbone._splits
if 'neck_split' in splits.keys():
encoder, neck, decoder = split_converter(
list_encdec, splits['backbone_split'],
splits['neck_split'])
else:
encoder, decoder = split_converter(list_encdec,
splits['backbone_split'])
neck = None
load_weights_backbone(model.backbone, encoder)
model.backbone.trainable = False
if self.task_config.darknet_load_decoder:
if neck != None:
load_weights_backbone(model.decoder.neck, neck)
model.decoder.neck.trainable = False
cfgheads = load_head(model.decoder.head, decoder)
model.decoder.head.trainable = False
load_weights_prediction_layers(cfgheads, model.head)
model.head.trainable = False
else:
"""Loading pretrained checkpoint."""
if not self.task_config.init_checkpoint:
return
ckpt_dir_or_file = self.task_config.init_checkpoint
if tf.io.gfile.isdir(ckpt_dir_or_file):
ckpt_dir_or_file = tf.train.latest_checkpoint(ckpt_dir_or_file)
# Restoring checkpoint.
if self.task_config.init_checkpoint_modules == 'all':
ckpt = tf.train.Checkpoint(**model.checkpoint_items)
status = ckpt.restore(ckpt_dir_or_file)
status.assert_consumed()
elif self.task_config.init_checkpoint_modules == 'backbone':
ckpt = tf.train.Checkpoint(backbone=model.backbone)
status = ckpt.restore(ckpt_dir_or_file)
status.expect_partial().assert_existing_objects_matched()
else:
assert "Only 'all' or 'backbone' can be used to initialize the model."
logging.info('Finished loading pretrained checkpoint from %s',
ckpt_dir_or_file)