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')