def create_dataset(opt, model_cfg):
data_config = parse_data_config(opt.data)
train_path = data_config["train"]
valid_path = data_config["valid"]
class_names = load_classes(data_config["names"])
train_loader, valid_loader, test_loader = None, None, None
train_dataset = ListDataset(train_path,
img_size=int(model_cfg[0]['width']),
augment=True,
multiscale=opt.multiscale_training)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=4,
pin_memory=True,
collate_fn=train_dataset.collate_fn,
)
#import pdb;pdb.set_trace()
valid_dataset = ListDataset(train_path,
img_size=int(model_cfg[0]['width']),
augment=False,
multiscale=False)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=opt.batch_size,
shuffle=False,
num_workers=4,
collate_fn=valid_dataset.collate_fn)
test_loader = valid_loader
return class_names, train_loader, valid_loader, test_loader
def _create_data_loader(img_path, batch_size, img_size, n_cpu, multiscale_training=False):
"""Creates a DataLoader for training.
:param img_path: Path to file containing all paths to training images.
:type img_path: str
:param batch_size: Size of each image batch
:type batch_size: int
:param img_size: Size of each image dimension for yolo
:type img_size: int
:param n_cpu: Number of cpu threads to use during batch generation
:type n_cpu: int
:param multiscale_training: Scale images to different sizes randomly
:type multiscale_training: bool
:return: Returns DataLoader
:rtype: DataLoader
"""
dataset = ListDataset(
img_path,
img_size=img_size,
multiscale=multiscale_training,
transform=AUGMENTATION_TRANSFORMS)
dataloader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=True,
num_workers=n_cpu,
pin_memory=True,
collate_fn=dataset.collate_fn)
return dataloader
def _create_validation_data_loader(img_path, batch_size, img_size, n_cpu):
"""
Creates a DataLoader for validation.
:param img_path: Path to file containing all paths to validation images.
:type img_path: str
:param batch_size: Size of each image batch
:type batch_size: int
:param img_size: Size of each image dimension for yolo
:type img_size: int
:param n_cpu: Number of cpu threads to use during batch generation
:type n_cpu: int
:return: Returns DataLoader
:rtype: DataLoader
"""
dataset = ListDataset(img_path,
img_size=img_size,
multiscale=False,
transform=DEFAULT_TRANSFORMS)
dataloader = DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=n_cpu,
pin_memory=True,
collate_fn=dataset.collate_fn)
return dataloader
def evaluate(model, path, iou_thres, conf_thres, nms_thres, img_size,
batch_size):
model.eval()
# Get dataloader
dataset = ListDataset(path,
img_size=img_size,
augment=False,
multiscale=False)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=0,
collate_fn=dataset.collate_fn)
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available(
) else torch.FloatTensor
labels = []
sample_metrics = [] # List of tuples (TP, confs, pred)
for batch_i, (_, _, imgs, targets) in enumerate(
tqdm.tqdm(dataloader, desc="Detecting objects")):
# Extract labels
if targets is None:
continue
labels += targets[:, 1].tolist()
# Rescale target to x1y1x2y2. YOLO outputs in XYWH and it gets converted in the IOU script later
targets[:, 2:] = xywh2xyxy(targets[:, 2:])
targets[:, 2:] *= img_size
imgs = Variable(imgs.type(Tensor), requires_grad=False)
with torch.no_grad():
outputs = model(imgs)
outputs = non_max_suppression(outputs,
conf_thres=conf_thres,
nms_thres=nms_thres)
sample_metrics += get_batch_statistics(outputs,
targets,
iou_threshold=iou_thres)
# Concatenate sample statistics
try:
true_positives, pred_scores, pred_labels = [
np.concatenate(x, 0) for x in list(zip(*sample_metrics))
]
precision, recall, AP, f1, ap_class = ap_per_class(
true_positives, pred_scores, pred_labels, labels)
except ValueError as error:
print('-----------------------------------------------')
print(error)
print('Model failed to detect any boxes in validation above threshold')
print('Zeros passed for all metrics')
print('-----------------------------------------------')
precision, recall, f1 = (None, None, None)
AP = np.array([0] * len(np.unique(labels)))
ap_class = np.unique(labels).astype("int32")
return precision, recall, AP, f1, ap_class
def evaluate(model, path, iou_thres, conf_thres, nms_thres, img_size, batch_size):
model.eval()
# Get dataloader
dataset = ListDataset(path, img_size=img_size, augment=False, multiscale=False)
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=False, num_workers=1, collate_fn=dataset.collate_fn
)
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
labels = []
sample_metrics = [] # List of tuples (TP, confs, pred)
for batch_i, (_, imgs, targets) in enumerate(tqdm.tqdm(dataloader, desc="Detecting objects")):
# Extract labels
labels += targets[:, 1].tolist()
# Rescale target
targets[:, 2:] = xywh2xyxy(targets[:, 2:])
targets[:, 2:] *= img_size
imgs = Variable(imgs.type(Tensor), requires_grad=False)
with torch.no_grad():
outputs = model(imgs)
outputs = non_max_suppression(outputs, conf_thres=conf_thres, nms_thres=nms_thres)
sample_metrics += get_batch_statistics(outputs, targets, iou_threshold=iou_thres)
# Concatenate sample statistics
true_positives, pred_scores, pred_labels = [np.concatenate(x, 0) for x in list(zip(*sample_metrics))]
precision, recall, AP, f1, ap_class = ap_per_class(true_positives, pred_scores, pred_labels, labels)
return precision, recall, AP, f1, ap_class
def main() -> None:
testset = ListDataset(args.img_list_file_path,
img_size=416,
multiscale=False,
transform=DEFAULT_TRANSFORMS)
# get_dataset_ap(testset, "test/barcode5/labels")
pres, recas = get_dataset_ap(testset, args.pred_label_path)
draw(pres, recas)
def evaluate(model, path, iou_thres, conf_thres, nms_thres, img_size,
batch_size):
model.eval()
# Get dataloader
dataset = ListDataset(path, img_size=img_size, augment=False)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=1,
collate_fn=dataset.collate_fn)
if TQDM_USE: dataloader = tqdm(dataloader)
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available(
) else torch.FloatTensor
labels = []
sample_metrics = [] # List of tuples (TP, confs, pred)
for batch_i, (_, imgs, targets) in enumerate(dataloader):
#imgs.shape(batch_size, 3, img_size, img_size), img_size is 416 by default
#targets.shape(num_bbox, 6), the 6 values are (idx, labels, x, y, w, h)
# Extract labels
labels += targets[:, 1].tolist()
# Rescale target
targets[:, 2:] = xywh2xyxy(
targets[:, 2:]) # convert the coordinates from xywh to xyxy
targets[:,
2:] *= img_size # scale the normalized targets to image size
imgs = Variable(imgs.type(Tensor), requires_grad=False)
with torch.no_grad():
#outputs.shape(batch_size, num_yolo*num_anchors*grid_size*grid_size, 85)
outputs = model(imgs)
# the outputs from NMS is a list, which has batch_size elements
# each element of outputs is the prediction bboxes of one sample
# each bbox inclues (x1, y1, x2, y2, object_conf, class_score, class_pred_label)
outputs = non_max_suppression(outputs,
conf_thres=conf_thres,
nms_thres=nms_thres)
sample_metrics += get_batch_statistics(outputs,
targets,
iou_threshold=iou_thres)
# Concatenate sample statistics
true_positives, pred_scores, pred_labels = [
np.concatenate(x, 0) for x in list(zip(*sample_metrics))
]
precision, recall, AP, f1, ap_class = ap_per_class(true_positives,
pred_scores,
pred_labels, labels)
return precision, recall, AP, f1, ap_class
def test_ListDataset(self):
img_paht = 'E:/ML_data/trainvalno5k.txt'
dataset = ListDataset(img_paht, augment=True, multiscale=True)
print(len(dataset))
dataloader_list = DataLoader(dataset,
batch_size=6,
shuffle=True,
num_workers=0,
pin_memory=True,
collate_fn=dataset.collate_fn)
def get_data(dataset_name, dataset_path, num_imgs, n_cpu):
'''Get data'''
if dataset_name == 'coco':
dataset = ListDataset(dataset_path)
elif dataset_name == 'nwpuvhr':
dataset = NwpuvhrDataset(dataset_path)
elif dataset_name == 'dota':
dataset = DotaDataset(dataset_path)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=num_imgs,
shuffle=True,
num_workers=n_cpu)
return dataset, dataloader
def evaluate(model, path, iou_thres, conf_thres, nms_thres, img_size,
batch_size, device):
model.eval()
# Get dataloader
dataset = ListDataset(path,
img_size=img_size,
augment=False,
multiscale=False)
dataloader = tc.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=1,
collate_fn=dataset.collate_fn)
labels = []
sample_metrics = [] # List of tuples (TP, confs, pred)
try:
tq = tqdm.tqdm(dataloader, desc='Detecting objects', ncols=100)
for _, imgs, targets in tq:
imgs = imgs.to(device)
# Extract labels
labels += targets[:, 1].tolist()
# Rescale target
targets[:, 2:] = xywh2xyxy(targets[:, 2:])
targets[:, 2:] *= img_size
with tc.no_grad():
outputs = model(imgs)
outputs = non_max_suppression(outputs,
conf_thres=conf_thres,
nms_thres=nms_thres)
sample_metrics += get_batch_statistics(outputs,
targets,
iou_threshold=iou_thres)
finally:
tq.close()
# Concatenate sample statistics
true_positives, pred_scores, pred_labels = [
np.concatenate(x, 0) for x in list(zip(*sample_metrics))
]
precision, recall, AP, f1, ap_class = ap_per_class(true_positives,
pred_scores,
pred_labels, labels)
return precision, recall, AP, f1, ap_class
class_names = load_classes(data_config["names"])
# Initiate model
model = Darknet(opt.model_def).to(device)
model.apply(weights_init_normal)
# If specified we start from checkpoint
if opt.pretrained_weights:
if opt.pretrained_weights.endswith(".pth"):
model.load_state_dict(torch.load(opt.pretrained_weights))
else:
model.load_darknet_weights(opt.pretrained_weights)
# Get dataloader
dataset = ListDataset(train_path,
augment=False,
multiscale=opt.multiscale_training)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=0, # opt.n_cpu,
pin_memory=True,
collate_fn=dataset.collate_fn)
optimizer = torch.optim.Adam(model.parameters())
metrics = [
"grid_size",
"loss",
"x",
_PROTOTXT_PATH = '%s/SSD_caffe_model/ASU_model/deploy_6.prototxt' % HOME_DIR
# Default input caffemodel path
_CAFFEMODEL_PATH = '%s/SSD_caffe_model/ASU_model/SSD300_6_iter_120000.caffemodel' % HOME_DIR
# Default output pytorch weights path
_PTHMODEL_PATH = '%s/SSD_pytorch_model/SSD300_6_iter_120000.pth' % HOME_DIR
torch_folder_dataloader = DataLoader(ImageFolder('data/samples', img_size=300),
batch_size=1,
shuffle=False,
num_workers=1,
pin_memory=True)
torch_list_dataloader = DataLoader(ListDataset('data/samples/listfile.txt',
img_size=300,
normalized_labels=False),
batch_size=1,
shuffle=False,
num_workers=1,
pin_memory=True)
def parse_caffemodel(caffemodel):
"""Parse a caffemodel
Inputs:
string of caffemodel path
Returns:
Parsed model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
os.makedirs("checkpoints", exist_ok=True)
# Get data configuration
data_config = parse_data_config(opt.data_config)
train_path = data_config["train"]
valid_path = data_config["valid"]
class_names = load_classes(data_config["names"])
# multi-sacle image size
img_min_size = opt.img_size - 3*32
img_max_size = opt.img_size + 3*32
img_cur_size = img_max_size if opt.multiscale_training else opt.img_size
# Get dataloader
dataset = ListDataset(train_path, augment=True,
img_size=img_cur_size)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.n_cpu,
pin_memory=True,
collate_fn=dataset.collate_fn,
)
# Initiate model
model = Darknet(opt.model_def).to(device)
model.apply(weights_init_normal)
# If specified we start from checkpoint
if opt.pretrained_weights:
def evaluate(
model,
path: str,
transform,
iou_thres: float = 0.5,
conf_thres: float = 0.5,
nms_thres: float = 0.5,
img_size: int = 416,
batch_size: int = 8,
):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.eval()
dataset = ListDataset(list_path=path,
transform=transform,
img_size=img_size,
num_samples=None)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
shuffle=False,
num_workers=4,
collate_fn=dataset.collate_fn,
)
labels = []
sample_metrics = [] # List of tuples (TP, confs, pred)
for batch_i, (_, imgs, targets) in enumerate(
tqdm.tqdm(dataloader, desc="Detecting objects")):
labels += targets[:, 1].tolist()
targets[:, 2:] = xywh2xyxy(targets[:, 2:])
targets[:, 2:] *= img_size
imgs = imgs.requires_grad_(False).to(device)
with torch.no_grad():
outputs = model(imgs)
outputs = non_max_suppression(prediction=outputs,
conf_thres=conf_thres,
nms_thres=nms_thres)
sample_metrics += get_batch_statistics(outputs=outputs,
targets=targets,
iou_threshold=iou_thres)
# Concatenate sample statistics
if sample_metrics:
true_positives, pred_scores, pred_labels = [
np.concatenate(x, 0) for x in list(zip(*sample_metrics))
]
else:
true_positives, pred_scores, pred_labels = [
np.zeros(1), np.zeros(1), np.zeros(1)
]
precision, recall, AP, f1, ap_class = ap_per_class(true_positives,
pred_scores,
pred_labels, labels)
return precision, recall, AP, f1, ap_class
valid_path = data_config["valid"]
class_names = load_classes(data_config["names"])
# Initiate model
model = Darknet(opt.model_def).to(device)
model.apply(weights_init_normal)
# If specified we start from checkpoint
if opt.pretrained_weights:
if opt.pretrained_weights.endswith(".pth"):
model.load_state_dict(torch.load(opt.pretrained_weights))
else:
model.load_darknet_weights(opt.pretrained_weights)
# Get dataloader
dataset = ListDataset(train_path, augment=True)
dataloader = DataLoaderX(
dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=0,
pin_memory=False,
collate_fn=dataset.collate_fn)
# dataloader = DataPrefetcher(
# torch.utils.data.DataLoader(
# dataset,
# batch_size=opt.batch_size,
# shuffle=False,
# num_workers=opt.n_cpu,
# pin_memory=True,
# Initiate model
model = Darknet(opt.model_def, label_size=opt.label_size).to(device)
model.apply(weights_init_normal)
# If specified we start from checkpoint
if opt.pretrained_weights:
if opt.pretrained_weights.endswith(".pth"):
model.load_state_dict(torch.load(opt.pretrained_weights))
else:
model.load_darknet_weights(opt.pretrained_weights)
# Get dataloader
dataset = ListDataset(
train_path,
label_size=opt.label_size,
augment=False,
multiscale=opt.multiscale_training,
)
dataloader = DataLoader(
dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.n_cpu,
pin_memory=True,
collate_fn=dataset.collate_fn,
)
optimizer = torch.optim.Adam(model.parameters())
metrics = [
"grid_size",
def main():
# Hyperparameters parser
parser = argparse.ArgumentParser()
parser.add_argument("--year", type=str, default='2012', help="used to select training set")
parser.add_argument("--set", type=str, default='train', help="used to select training set")
parser.add_argument("--epochs", type=int, default=201, help="number of epochs")
parser.add_argument("--batch_size", type=int, default=8, help="size of each image batch")
parser.add_argument("--model_def", type=str, default="config/net/resnet_dropout.cfg", help="path to model definition file")
# parser.add_argument("--model_def", type=str, default="config/net/dqnyolo_large.cfg", help="path to model definition file")
# parser.add_argument("--model_def", type=str, default="config/net/dqnyolo_mini.cfg", help="path to model definition file")
# parser.add_argument("--model_def", type=str, default="config/net/dqnyolo_tiny.cfg", help="path to model definition file")
parser.add_argument("--img_size", type=int, default=416, help="size of each image dimension")
parser.add_argument("--opt_lr", type=float, default=1e-5, help="learning rate for optimizer")
parser.add_argument("--use_gpu", default=True, help="use GPU to accelerate training")
parser.add_argument("--shuffle_train", default=True, help="shuffle the training dataset")
parser.add_argument("--checkpoint_interval", type=int, default=20, help="interval between saving model weights")
parser.add_argument("--evaluation_interval", type=int, default=10, help="interval evaluations on validation set")
# parser.add_argument("--pretrained_weights", type=str, default="data/backbone/darknet53.conv.74", help="if specified starts from checkpoint model")
# parser.add_argument("--pretrained_weights", type=str, default="logs/model/model_params_200.ckpt", help="if specified starts from checkpoint model")
parser.add_argument("--pretrained_weights", default=False, help="if specified starts from checkpoint model")
opt = parser.parse_args()
print(opt)
if opt.use_gpu is True:
if torch.cuda.is_available():
device = torch.device('cuda')
else:
raise RuntimeError("Current Torch doesn't have GPU support.")
else:
device = torch.device('cpu')
logger = SummaryWriter(exist_or_create_folder("./logs/tb/"))
# Initiate model
eval_model = Darknet(opt.model_def).to(device)
if opt.pretrained_weights:
print("Initialize model with pretrained_model")
if opt.pretrained_weights.endswith(".ckpt"):
eval_model.load_state_dict(torch.load(opt.pretrained_weights))
else:
eval_model.load_darknet_weights(opt.pretrained_weights)
else:
print("Initialize model randomly")
eval_model.apply(weights_init_normal)
# eval_model.load_state_dict(torch.load("./logs/saved_exp/master-v2/model_params_80.ckpt"))
print(eval_model)
summary(eval_model, (3, 416, 416))
learn_batch_counter = 0 # for logger update (total numbers)
batch_size = opt.batch_size
# Get dataloader
print("Begin loading train dataset ......")
t_load_data = time.time()
dataset = torchvision.datasets.VOCDetection(root='data/VOC/',
year=opt.year,
image_set=opt.set,
transforms=None,
download=True)
dataset_dict = trans_voc(dataset)
dataset = ListDataset(dataset_dict)
loader = torch.utils.data.DataLoader(
dataset,
batch_size=opt.batch_size,
shuffle=opt.shuffle_train,
pin_memory=True,
collate_fn=dataset.collate_fn,
)
print("Complete loading train dataset in {} s".format(time.time() - t_load_data))
optimizer = torch.optim.Adam(eval_model.parameters(), lr=opt.opt_lr)
# Warmup and learning rate decay
scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, opt.epochs)
# 5 epoch warmup, lr from 1e-5 to 1e-4, after that schedule as after_scheduler
scheduler_warmup = GradualWarmupScheduler(optimizer, multiplier=10, total_epoch=10,
after_scheduler=scheduler_cosine)
start_time = time.time()
for i_epoch in range(opt.epochs):
eval_model.train()
for i_batch, (_, imgs, raw_targets, transform_params, tar_boxes) in enumerate(loader):
print("\n++++++++++ i_epoch-i_batch {}-{} ++++++++++".format(i_epoch, i_batch))
batch_step_counter = 0
if len(imgs) != batch_size:
print("Current batch size is smaller than opt.batch_size!")
continue
imgs = imgs.to(device)
raw_targets = raw_targets.to(device)
tar_boxes = tar_boxes.to(device)
input_img = imgs
if i_epoch == 0 and i_batch == 0:
logger.add_graph(eval_model, input_img)
# print(raw_targets)
# print(raw_targets.size())
# print(raw_targets[:, :, :, 6:].size())
# print(raw_targets[:, :, :, 0].unsqueeze(3).size())
cls_targets = torch.cat((raw_targets[:, :, :, 0].unsqueeze(3), raw_targets[:, :, :, 6:]), 3)
# print(cls_targets.size())
loss, pred = eval_model(input_img, cls_targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
batch_step_counter += 1
learn_batch_counter += 1
print("Ep-bt: {}-{} | Loss: {}".format(i_epoch, i_batch, loss.item()))
logger.add_scalar('loss/loss', loss.item(), learn_batch_counter)
if (i_epoch + 1) % opt.checkpoint_interval == 0:
print("Saving model in epoch {}".format(i_epoch))
torch.save(eval_model.state_dict(),
exist_or_create_folder("./logs/model/model_params_{}.ckpt".format(i_epoch)))
# Evaluate the model on the validation set
if (i_epoch + 1) % opt.evaluation_interval == 0:
precision, recall, AP, f1, ap_class = evaluate(
eval_model,
[opt.year, 'val'],
[0.5, 0.5, 0.5],
batch_size,
True,
diagnosis_code=1
)
evaluation_metrics = [
("val_precision", precision.mean()),
("val_recall", recall.mean()),
("val_mAP", AP.mean()),
("val_f1", f1.mean()),
]
for tag, value in evaluation_metrics:
logger.add_scalar("val/{}".format(tag), value.item(), i_epoch)
# Print class APs and mAP
ap_table = [["Index", "Class name", "AP"]]
for i, c in enumerate(ap_class):
ap_table += [[c, val2labels(c), "%.5f" % AP[i]]]
print(AsciiTable(ap_table).table)
print(f"---- validation mAP {AP.mean()}")
# Evaluate the model on the training set
if (i_epoch + 1) % opt.evaluation_interval == 0:
precision, recall, AP, f1, ap_class = evaluate(
eval_model,
[opt.year, 'train'],
[0.5, 0.5, 0.5],
batch_size,
True,
diagnosis_code=1
)
evaluation_metrics = [
("train_precision", precision.mean()),
("train_recall", recall.mean()),
("train_mAP", AP.mean()),
("train_f1", f1.mean()),
]
for tag, value in evaluation_metrics:
logger.add_scalar("train/{}".format(tag), value.item(), i_epoch)
# Print class APs and mAP
ap_table = [["Index", "Class name", "AP"]]
for i, c in enumerate(ap_class):
ap_table += [[c, val2labels(c), "%.5f" % AP[i]]]
print(AsciiTable(ap_table).table)
print(f"---- training mAP {AP.mean()}")
# Warmup and lr decay
scheduler_warmup.step()
# Free GPU memory
torch.cuda.empty_cache()
total_train_time = time.time() - start_time
print("Training complete in {} hours".format(total_train_time / 3600))
def evaluate(model,
dataset_des,
thres,
batch_size,
init_dim_cluster,
diagnosis_code=0):
model.eval()
iou_thres, conf_thres, nms_thres = thres
# Get dataloader
print("Begin loading validation dataset ......")
t_load_data = time.time()
dataset = torchvision.datasets.VOCDetection(root='data/VOC/',
year=dataset_des[0],
image_set=dataset_des[1],
transforms=None,
download=True)
dataset_dict = trans_voc(dataset)
dataset = ListDataset(dataset_dict)
loader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
shuffle=False,
pin_memory=True,
collate_fn=dataset.collate_fn,
)
print("Complete loading validation dataset in {} s".format(time.time() -
t_load_data))
labels = []
sample_metrics = [] # List of tuples (TP, confs, pred)
for i_batch, (img_ind, img, raw_targets, transform_params,
tar_boxes) in enumerate(loader):
print("\n++++++++++ i_batch (val) {} ++++++++++".format(i_batch))
batch_step_counter = 0
# Extract labels: raw_targets -- [t_conf, t_cls, x, y, w, h, one_hot_t_cls(20)]
labels += tar_boxes[:, 1].tolist()
if len(img) != batch_size:
print("Current batch size is smaller than opt.batch_size!")
continue
img = img.to('cuda')
raw_targets = raw_targets.to('cuda')
tar_boxes = tar_boxes.to('cuda')
input_img = img
with torch.no_grad():
pred_conf_cls = model(input_img)
pred_conf_cls = pred_conf_cls.permute(0, 2, 3, 1)
pred_conf = torch.sigmoid(pred_conf_cls[:, :, :, 0])
pred_cls = torch.sigmoid(pred_conf_cls[:, :, :, 1:])
obj_mask = pred_conf > conf_thres
obj_mask = obj_mask.byte().to('cuda')
if diagnosis_code == 0:
pass
if diagnosis_code == 1:
# localization ground-truth
pred_bbox = raw_targets[:, :, :, 2:6]
# pred_conf_cls = pred_conf_cls.permute(0, 2, 3, 1)
# pred_conf = torch.sigmoid(pred_conf_cls[:, :, :, 0])
# pred_cls = torch.sigmoid(pred_conf_cls[:, :, :, 1:])
if diagnosis_code == 2:
# classification ground-truth
pass
if diagnosis_code == 3:
# full ground-truth
pred_bbox = raw_targets[:, :, :, 2:6]
pred_conf = raw_targets[:, :, :, 0]
pred_cls = raw_targets[:, :, :, 6:]
# if i_batch < 20:
# im = Image.open("data/VOC/VOCdevkit/VOC2007/JPEGImages/{}".format(img_ind[0])).convert('RGB')
# plot_detection_result(im, img_ind, pred_bbox, pred_conf, pred_cls, transform_params, iou=0.9)
pred_outputs = torch.cat((pred_bbox, pred_conf.unsqueeze(3), pred_cls),
dim=3)
b, w, h, d = pred_outputs.shape
pred_outputs = pred_outputs.view(b, w * h, d)
outputs = non_max_suppression(pred_outputs,
conf_thres=conf_thres,
nms_thres=nms_thres)
sample_metrics += get_batch_statistics(outputs,
tar_boxes,
iou_threshold=iou_thres)
# Concatenate sample statistics
true_positives, pred_scores, pred_labels = [
np.concatenate(x, 0) for x in list(zip(*sample_metrics))
]
precision, recall, AP, f1, ap_class = ap_per_class(true_positives,
pred_scores,
pred_labels, labels)
return precision, recall, AP, f1, ap_class
darknet_pretrained = False
if opt.pretrained_weights:
print("load pretrained weights")
if opt.pretrained_weights.endswith(".pth"):
"""Load darknet weights"""
model.load_state_dict(torch.load(opt.pretrained_weights))
darknet_pretrained = True
else:
"""Load checkpoint weights"""
model.load_darknet_weights(opt.pretrained_weights)
train_transform, test_transform = get_transform(
img_size=opt.img_size, darknet_pretrained=darknet_pretrained)
dataset = ListDataset(
list_path=train_path,
transform=train_transform,
img_size=opt.img_size,
num_samples=None,
)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.n_cpu,
pin_memory=True,
collate_fn=dataset.collate_fn,
)
optimizer = torch.optim.Adam(params=model.parameters(), lr=0.005)
metrics, formats = get_constants()
for epoch in range(opt.epochs):
valid_path = settings.valid
class_names = load_classes(settings.names)
# Initiate model
model = Darknet(settings.model_def).to(device)
model.apply(weights_init_normal)
# If specified we start from checkpoint
if settings.pretrained_weights:
if settings.pretrained_weights.endswith(".pth"):
model.load_state_dict(torch.load(settings.pretrained_weights))
else:
model.load_darknet_weights(settings.pretrained_weights)
# Get dataloader
dataset = ListDataset(settings.dataset_prefix, train_path, augment=True, multiscale=settings.multiscale_training)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=settings.batch_size,
shuffle=True,
num_workers=settings.n_cpu,
pin_memory=True,
collate_fn=dataset.collate_fn,
)
valid_dataset = ListDataset(settings.dataset_prefix, valid_path, img_size=settings.img_size, augment=False, multiscale=False)
valid_dataloader = torch.utils.data.DataLoader(
valid_dataset, batch_size=settings.batch_size, shuffle=False, num_workers=1, collate_fn=valid_dataset.collate_fn
)
optimizer = torch.optim.Adam(model.parameters(), lr=settings.lr)