from core.loss import JointsMSELoss
from core.make_dataset import CocoDataset
from core.make_ground_truth import GroundTruth
from core.metric import PCK
from utils.work_flow import get_model, print_model_summary
from configuration.base_config import Config
from utils.tools import get_config_params
if __name__ == '__main__':
# GPU settings
gpus = tf.config.list_physical_devices("GPU")
if gpus:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
cfg = get_config_params(Config.TRAINING_CONFIG_NAME)
hrnet = get_model(cfg)
print_model_summary(hrnet)
# Dataset
coco = CocoDataset(config_params=cfg, dataset_type="train")
dataset, dataset_length = coco.generate_dataset()
# loss and optimizer
loss = JointsMSELoss()
optimizer = tf.optimizers.Adam(learning_rate=1e-3)
# metircs
loss_metric = tf.metrics.Mean()
pck = PCK()
accuracy_metric = tf.metrics.Mean()
def print_model_summary(network):
config_params = get_config_params(Config.TRAINING_CONFIG_NAME)
network.build(input_shape=(None, config_params.IMAGE_HEIGHT,
config_params.IMAGE_WIDTH,
config_params.CHANNELS))
network.summary()