def main(data_config, model_def, weights_path, batch_size, img_size, iou_thres,
conf_thres, nms_thres, nb_cpu):
data_config = parse_data_config(update_path(data_config))
valid_path = update_path(data_config["valid"])
class_names = load_classes(update_path(data_config["names"]))
# Initiate model
model = Darknet(model_def).to(DEVICE)
if weights_path.endswith(".weights"):
# Load darknet weights
model.load_darknet_weights(weights_path)
else:
# Load checkpoint weights
model.load_state_dict(torch.load(weights_path))
print("Compute mAP...")
precision, recall, AP, f1, ap_class = evaluate_model(
model,
path_data=valid_path,
iou_thres=iou_thres,
conf_thres=conf_thres,
nms_thres=nms_thres,
img_size=img_size,
batch_size=batch_size,
nb_cpu=nb_cpu,
)
print("Average Precisions:")
for i, c in enumerate(ap_class):
print(f"+ Class '{c}' ({class_names[c]}) - AP: {AP[i]}")
print(f"mAP: {AP.mean()}")
def __getitem__(self, idx):
# ---------
# Image
# ---------
idx = idx % len(self.img_files)
img_path = self.img_files[idx].rstrip()
img_path = update_path(img_path)
# Extract image as PyTorch tensor
img = transforms.ToTensor()(Image.open(img_path).convert('RGB'))
# Handle images with less than three channels
if len(img.shape) != 3:
img = img.unsqueeze(0)
img = img.expand((3, img.shape[1:]))
_, h, w = img.shape
h_factor, w_factor = (h, w) if self.normalized_labels else (1, 1)
# Pad to square resolution
img, pad = pad_to_square(img, 0)
_, padded_h, padded_w = img.shape
# ---------
# Label
# ---------
label_path = self.label_files[idx].rstrip()
label_path = update_path(label_path)
targets = None
if os.path.exists(label_path):
boxes = torch.from_numpy(np.loadtxt(label_path).reshape(-1, 5))
# Extract coordinates for unpadded + unscaled image
x1 = w_factor * (boxes[:, 1] - boxes[:, 3] / 2)
y1 = h_factor * (boxes[:, 2] - boxes[:, 4] / 2)
x2 = w_factor * (boxes[:, 1] + boxes[:, 3] / 2)
y2 = h_factor * (boxes[:, 2] + boxes[:, 4] / 2)
# Adjust for added padding
x1 += pad[0]
y1 += pad[2]
x2 += pad[1]
y2 += pad[3]
# Returns (x, y, w, h)
boxes[:, 1] = ((x1 + x2) / 2) / padded_w
boxes[:, 2] = ((y1 + y2) / 2) / padded_h
boxes[:, 3] *= w_factor / padded_w
boxes[:, 4] *= h_factor / padded_h
targets = torch.zeros((len(boxes), 6))
targets[:, 1:] = boxes
# Apply augmentations
if self.augment.get('hflip', False) and np.random.random() < 0.5:
img, targets = horisontal_flip(img, targets)
if self.augment.get('vflip', False) and np.random.random() < 0.5:
img, targets = vertical_flip(img, targets)
return img_path, img, targets
def __init__(self, list_path, img_size=416, augment=None, normalized_labels=True, max_objects=100):
"""Dataset defined by list of images/annot
:param list list_path:
:param int img_size:
:param dict augment: options like {'scaling': 0.1, 'vflip': None, 'hflip': None}
:param bool normalized_labels:
:param int max_objects:
"""
with open(list_path, "r") as file:
img_files = [update_path(fn.strip()) for fn in file.readlines()]
# get the labels for each image if image exista
label_files = [
os.path.join(os.path.dirname(imp).replace("images", "labels"),
os.path.splitext(os.path.basename(imp))[0] + ".txt")
for imp in img_files if os.path.isfile(imp)
]
logging.debug("From %i listed, found %i images", len(img_files), len(label_files))
# filter existing image and annotation
path_img_lbs = [(p_img, p_lbs) for p_img, p_lbs in zip(img_files, label_files)
if os.path.isfile(p_img) and os.path.isfile(p_lbs)]
assert path_img_lbs, 'missing images/annotations'
self.img_files, self.label_files = list(zip(*path_img_lbs))
logging.debug("From %i listed found %i annotation", len(path_img_lbs), len(self.label_files))
self.augment = augment if augment else {}
self.img_size = img_size
self.max_objects = max_objects
self.normalized_labels = normalized_labels
self.batch_count = 0
scaling = self.augment.get('scaling', 0.)
self.min_size = round(self.img_size * (1. - scaling) / 32) * 32
self.max_size = round(self.img_size * (1. + scaling) / 32) * 32
def __getitem__(self, idx):
# ---------
# Image
# ---------
idx = idx % len(self.img_files)
img_path = self.img_files[idx].rstrip()
img_path = update_path(img_path)
# Extract image as PyTorch tensor
img = transforms.ToTensor()(Image.open(img_path).convert('RGB'))
# Handle images with less than three channels
if len(img.shape) != 3:
img = img.unsqueeze(0)
img = img.expand((3, img.shape[1:]))
_, h, w = img.shape
h_factor, w_factor = (h, w) if self.normalized_labels else (1, 1)
# Pad to square resolution
img, pad = pad_to_square(img, 0)
_, padded_h, padded_w = img.shape
# ---------
# Label
# ---------
label_path = self.label_files[idx].rstrip()
label_path = update_path(label_path)
targets = None
if os.path.exists(label_path):
with warnings.catch_warnings(record=True) as w:
data = np.loadtxt(label_path)
# in case the file is empty
if len(data) > 0:
boxes = torch.from_numpy(data.reshape(-1, 5))
targets = ListDataset._convert_boxes_to_targets(
boxes, pad, (w_factor, h_factor), (padded_w, padded_h))
# Apply augmentations
if self.augment.get('hflip', False) and np.random.random() < 0.5:
img, targets = horisontal_flip(img, targets)
if self.augment.get('vflip', False) and np.random.random() < 0.5:
img, targets = vertical_flip(img, targets)
return img_path, img, targets
def main(data_config, model_def, trained_weights, augment, multiscale,
img_size, grad_accums, evaluation_interval, checkpoint_interval,
batch_size, epochs, path_output, nb_cpu, amp):
path_output = update_path(path_output)
os.makedirs(path_output, exist_ok=True)
shutil.copy(data_config, os.path.join(path_output, os.path.basename(data_config)))
_, _, _, local_vars = inspect.getargvalues(inspect.currentframe())
args = {arg: local_vars[arg] for arg in inspect.getfullargspec(main).args}
with open(os.path.join(path_output, 'script-config.yaml'), 'w') as fp:
yaml.dump(dict(args), fp)
logger = Logger(os.path.join(path_output, "logs"))
# Get data configuration
data_config = parse_data_config(update_path(data_config))
train_path = update_path(data_config["train"])
valid_path = update_path(data_config["valid"])
class_names = load_classes(update_path(data_config["names"]))
# Initiate model
assert os.path.isfile(model_def), 'missing: %s' % model_def
model = Darknet(update_path(model_def)).to(DEVICE)
model.apply(weights_init_normal)
# If specified we start from checkpoint
if trained_weights:
assert os.path.isfile(trained_weights), 'missing: %s' % trained_weights
if trained_weights.endswith(".pth"):
model.load_state_dict(torch.load(trained_weights))
else:
model.load_darknet_weights(trained_weights)
augment = dict(zip(augment, [True] * len(augment))) if augment else {}
augment["scaling"] = multiscale
# Get dataloader
assert os.path.isfile(train_path), 'missing: %s' % train_path
dataset = ListDataset(train_path, augment=augment, img_size=img_size)
dataloader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=True,
num_workers=nb_cpu,
pin_memory=True,
collate_fn=dataset.collate_fn,
)
for param in model.parameters():
param.requires_grad = True
# optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
optimizer = torch_optimizer.RAdam(model.parameters(), lr=0.0001)
# optimizer = torch_optimizer.Yogi(model.parameters(), lr=0.0001)
# optimizer = torch_optimizer.AdaBound(model.parameters(), lr=0.0001)
# Creates once at the beginning of training
scaler = torch.cuda.amp.GradScaler() if amp else None
for epoch in tqdm.tqdm(range(epochs), desc='Training epoch'):
model.train()
# start_time = time.time()
pbar_batch = tqdm.tqdm(total=len(dataloader))
train_metrics = []
for batch_i, (_, imgs, targets) in enumerate(dataloader):
model, batch_metric = training_batch(dataloader, model, optimizer, epochs,
epoch, batch_i, imgs, targets, grad_accums, img_size, scaler)
loss = batch_metric['loss']
pbar_batch.set_description("training batch loss=%.5f" % loss)
train_metrics.append(batch_metric)
pbar_batch.update()
pbar_batch.close()
pbar_batch.clear()
train_metrics_all = {m: [] for m in train_metrics[0]}
_ = [train_metrics_all[k].append(bm[k]) for bm in train_metrics for k in bm]
logger.list_scalars_summary([(k, np.mean(train_metrics_all[k]))
for k in train_metrics_all], step=epoch, phase='train')
if epoch % evaluation_interval == 0:
assert os.path.isfile(valid_path), 'missing: %s' % valid_path
evaluate_epoch(model, valid_path, img_size, batch_size, epoch, class_names, logger, nb_cpu)
if epoch % checkpoint_interval == 0:
torch.save(model.state_dict(), os.path.join(path_output, "yolov3_ckpt_%05d.pth" % epoch))
def main(data_config, model_def, trained_weights, augment, multiscale,
img_size, grad_accums, evaluation_interval, checkpoint_interval,
batch_size, epochs, path_output, nb_cpu):
path_output = update_path(path_output)
os.makedirs(path_output, exist_ok=True)
logger = Logger(os.path.join(path_output, "logs"))
# Get data configuration
data_config = parse_data_config(update_path(data_config))
train_path = update_path(data_config["train"])
valid_path = update_path(data_config["valid"])
class_names = load_classes(update_path(data_config["names"]))
# Initiate model
assert os.path.isfile(model_def)
model = Darknet(update_path(model_def)).to(DEVICE)
model.apply(weights_init_normal)
# If specified we start from checkpoint
if trained_weights:
assert os.path.isfile(trained_weights)
if trained_weights.endswith(".pth"):
model.load_state_dict(torch.load(trained_weights))
else:
model.load_darknet_weights(trained_weights)
augment = dict(zip(augment, [True] * len(augment))) if augment else {}
augment["scaling"] = multiscale
# Get dataloader
assert os.path.isfile(train_path)
dataset = ListDataset(train_path, augment=augment, img_size=img_size)
dataloader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=True,
num_workers=nb_cpu,
pin_memory=True,
collate_fn=dataset.collate_fn,
)
optimizer = torch.optim.Adam(model.parameters())
for epoch in tqdm.tqdm(range(epochs), desc='Training epoch'):
model.train()
# start_time = time.time()
pbar_batch = tqdm.tqdm(total=len(dataloader))
train_metrics = []
for batch_i, (_, imgs, targets) in enumerate(dataloader):
model, batch_metric = training_batch(dataloader, model, optimizer, epochs,
epoch, batch_i, imgs, targets, grad_accums, logger, img_size)
loss = batch_metric['loss']
pbar_batch.set_description("training batch loss=%.5f" % loss)
train_metrics.append(batch_metric)
pbar_batch.update()
pbar_batch.close()
pbar_batch.clear()
if epoch % evaluation_interval == 0:
assert os.path.isfile(valid_path)
evaluate_epoch(model, valid_path, img_size, batch_size, epoch, class_names, logger, nb_cpu)
if epoch % checkpoint_interval == 0:
torch.save(model.state_dict(), os.path.join(path_output, "yolov3_ckpt_%05d.pth" % epoch))
train_metrics_all = {m: [] for m in train_metrics[0]}
_ = [train_metrics_all[k].append(bm[k]) for bm in train_metrics for k in bm]
logger.list_scalars_summary([(k, np.mean(train_metrics_all[k]))
for k in train_metrics_all], step=epoch, phase='train')