def prettyShow(recovered_sample): image = draw_bounding_boxes(recovered_sample['image'], recovered_sample['bounding_boxes']) print( '$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$' ) print('ANGLE: ') print(recovered_sample['angle']) print( '$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$' ) image.show()
def draw_anchors(*, config='config', nb=100, only_groundtruth=False): cfg = _read_config(config) anchor_list = _build_anchor_list(cfg) (dataset, _), _ = _build_dataset(cfg, anchor_list) bgid = dataset.class_to_idx['background'] mean = cfg['mean'] std = cfg['std'] image_size = cfg['image_size'] for i in range(nb): x, gt_boxes, encoding = dataset[i] gt_boxes = [ ((x*image_size, y*image_size, w*image_size, h*image_size), dataset.idx_to_class[class_id]) for ((x, y, w, h), class_id) in gt_boxes ] x = x.numpy() x = x.transpose((1, 2, 0)) x = x * np.array(std) + np.array(mean) x = x.astype('float32') anchor_boxes = [] if not only_groundtruth: for j, e in enumerate(encoding): B, C = e A = anchor_list[j] hcoords, wcoords, k_id = np.where(C != bgid) for h, w, k in zip(hcoords, wcoords, k_id): a = image_size * A[h, w, k] a = a.tolist() c = C[h, w, k] c = dataset.idx_to_class[c] anchor_boxes.append((a, c)) pad = 100 im = np.zeros((x.shape[0] + pad * 2, x.shape[1] + pad * 2, x.shape[2])) im[pad:-pad, pad:-pad] = x im = draw_bounding_boxes(im, gt_boxes, color=(1, 0, 0), text_color=(1, 0, 0), pad=pad) im = draw_bounding_boxes(im, anchor_boxes, color=(0, 1, 0), text_color=(0, 1, 0), pad=pad) imsave(os.path.join(cfg['out_folder'], 'anchors', 'sample_{:05d}.jpg'.format(i)), im)
def test( *, in_folder='test_images', out_folder='test_results', model='out/model.th', score_threshold=None, topk=10, iou_threshold=None, background_threshold=0.5, use_nms=True, out=None, cuda=False): if not os.path.exists(out_folder): os.makedirs(out_folder) model = torch.load(model, map_location=lambda storage, loc: storage) if cuda: model = model.cuda() model.eval() filenames = [os.path.join(in_folder, f) for f in os.listdir(in_folder)] for filename in filenames: t0 = time.time() try: im = default_loader(filename) except OSError: continue w, h = im.size x = model.transform(im) scale_w, scale_h = w / x.size(2), h / x.size(1) X = x.view(1, x.size(0), x.size(1), x.size(2)) if cuda: X = X.cuda() X = Variable(X) Ypred = model(X) BP = [ bp.data.cpu(). view(bp.size(0), -1, model.num_coords, bp.size(2), bp.size(3)). permute(0, 3, 4, 1, 2). numpy() for bp, cp in Ypred] CP = [ cp.data.cpu(). view(cp.size(0), -1, model.nb_classes, cp.size(2), cp.size(3)). permute(0, 3, 4, 1, 2). numpy() for bp, cp in Ypred] X = X.data.cpu().numpy() x = X[0] x = x.transpose((1, 2, 0)) x = x * np.array(model.std) + np.array(model.mean) x = x.astype('float32') pred_boxes = [] for j in range(len(Ypred)): cp = CP[j][0]#cl1_score,cl2_score,... bp = BP[j][0]#4 if model.use_discrete_coords: bp = _get_coords(bp, model.coords_discretization) A = model.anchor_list[j] cp = get_probas(cp, model.classif_loss_name, axis=3) boxes = decode_bounding_box_list( bp, cp, A, image_size=model.image_size, include_scores=True, variance=model.config['variance'] ) pred_boxes.extend(boxes) if score_threshold is None: score_threshold = model.config['nms_score_threshold'] else: score_threshold = float(score_threshold) if iou_threshold is None: iou_threshold = model.config['nms_iou_threshold'] else: iou_threshold = float(iou_threshold) bgid = model.class_to_idx['background'] idx_to_class = {i: c for c, i in model.class_to_idx.items()} if use_nms: pred_boxes = non_maximal_suppression_per_class( pred_boxes, iou_threshold=iou_threshold, background_class_id=bgid, score_threshold=score_threshold ) else: def get_box(box, scores): bg_score = scores[0] if bg_score >= background_threshold: return box, 0, bg_score else: cl= 1 + scores[1:].argmax() score = scores[1:].max() return box, cl, score pred_boxes = [get_box(box, scores) for box, scores in pred_boxes] pred_boxes = sorted(pred_boxes, key=lambda p:p[2], reverse=True) pred_boxes = [(box, class_id, score) for box, class_id, score in pred_boxes if class_id != bgid] pred_boxes = [(box, idx_to_class[class_id], score) for box, class_id, score in pred_boxes] pred_boxes = pred_boxes[0:topk] pred_boxes = [ ((x * scale_w, y * scale_h, w * scale_w, h * scale_h), class_name, score) for (x, y, w, h), class_name, score in pred_boxes ] for box, class_name, score in pred_boxes: print(class_name, score) im = np.array(im) pad = int(model.config['pad'] * scale_w) impadded = np.zeros((im.shape[0] + pad * 2, im.shape[1] + pad * 2, im.shape[2])) impadded = impadded.astype('uint8') impadded[pad:-pad, pad:-pad] = im im = draw_bounding_boxes(impadded, pred_boxes, color=(0, 255, 0), text_color=(0, 255, 0), pad=pad) delta = time.time() - t0 outf = os.path.join(out_folder, os.path.basename(filename)) print('Processed {} in {:.3f}s'.format(outf, delta)) imsave(outf, impadded)
def _evaluate_model(model, train, valid): print('Evaluation') bgid = train.dataset.class_to_idx['background'] t0 = time.time() model.eval() metrics = defaultdict(list) cfg = model.config class_ids = list(set(train.dataset.class_to_idx.values()) - set([bgid])) class_ids = sorted(class_ids) for split_name, loader in (('train', train), ('valid', valid)): t0 = time.time() im_index = 0 for batch, samples, in enumerate(loader): tt0 = time.time() X, (Y, Ypred), (bbox_true, bbox_pred), (class_true, class_pred) = _predict(model, samples) X = X.data.cpu().numpy() B = [[b for b, c in y] for y in Y] B = [torch.from_numpy(np.array(b)).float() for b in B] C = [[c for b, c in y] for y in Y] C = [torch.from_numpy(np.array(c)).long() for c in C] # B contains groundtruth bounding boxes for each scale # C contains groundtruth classes for each scale BP = [ bp.data.cpu().view(bp.size(0), -1, model.num_coords, bp.size(2), bp.size(3)).permute(0, 3, 4, 1, 2).numpy() for bp, cp in Ypred] CP = [ cp.data.cpu().view(cp.size(0), -1, model.nb_classes, cp.size(2), cp.size(3)).permute(0, 3, 4, 1, 2).numpy() for bp, cp in Ypred] # BP contains predicted bounding boxes for each scale # CP contains predicted classes for each scale for i in range(len(X)): gt_boxes = [] pred_boxes = [] gt_boxes_per_class = defaultdict(list) pred_boxes_per_class = defaultdict(list) # for each example i in mini-batch x = X[i] x = x.transpose((1, 2, 0)) x = x * np.array(cfg['std']) + np.array(cfg['mean']) x = x.astype('float32') for j in range(len(Y)): # for each scale j ct = C[j][i]# class_id cp = CP[j][i]#cl1_score,cl2_score,... bt = B[j][i]#4 bp = BP[j][i]#4 if cfg['use_discrete_coords']: bp = _get_coords(bp, model.coords_discretization) A = model.anchor_list[j] # get groundtruth boxes gt_boxes.extend(decode_bounding_box_list( bt, ct, A, include_scores=False, image_size=cfg['image_size'], variance=cfg['variance'], )) # get predicted boxes cp = get_probas(cp, cfg['classif_loss'], axis=3) pred_boxes.extend(decode_bounding_box_list( bp, cp, A, include_scores=True, image_size=cfg['image_size'], variance=cfg['variance'] )) gt_boxes = [(box, class_id) for box, class_id in gt_boxes if class_id != bgid] for class_id in class_ids: gt_boxes_per_class[class_id].extend([ box for box, box_class_id in gt_boxes if class_id == box_class_id] ) # use the predicted boxes and groundtruth boxes to compute precision and recall # PER image pred_boxes = non_maximal_suppression_per_class( pred_boxes, background_class_id=bgid, iou_threshold=cfg['nms_iou_threshold'], score_threshold=cfg['nms_score_threshold'], ) for box, class_id, score in pred_boxes: pred_boxes_per_class[class_id].append((box, score)) pred_boxes = pred_boxes[0:cfg['nms_topk']] P = [] R = [] for class_id in class_ids: t = [box for box, cl in gt_boxes if cl == class_id] if len(t) == 0: continue p = [box for box, cl, score in pred_boxes if cl == class_id] prec = precision(p, t, iou_threshold=cfg['eval_iou_threshold']) re = recall(p, t, iou_threshold=cfg['eval_iou_threshold']) metrics['precision_' + train.dataset.idx_to_class[class_id] + '_' + split_name].append(prec) metrics['recall_' + train.dataset.idx_to_class[class_id] + '_' + split_name].append(re) P.append(prec) R.append(re) metrics['precision_' + split_name].append(np.mean(P)) metrics['recall_' + split_name].append(np.mean(R)) gt_boxes = [(box, train.dataset.idx_to_class[class_id]) for box, class_id in gt_boxes] pred_boxes = [(box, train.dataset.idx_to_class[class_id], score) for box, class_id, score in pred_boxes] # compute mAP and AP PER image recalls_mAP = np.linspace(0, 1, 11) AP_for_recall = defaultdict(list) AP_for_class = [] for class_id in class_ids: APs = average_precision( pred_boxes_per_class[class_id], gt_boxes_per_class[class_id], iou_threshold=cfg['eval_iou_threshold'], recalls_mAP=recalls_mAP, aggregate=False, ) if APs is None: continue AP = np.mean(APs) AP_for_class.append(AP) metrics['AP_' + train.dataset.idx_to_class[class_id] + '_' + split_name].append(AP) for r, ap in zip(recalls_mAP, APs): m = 'AP(rec_{:.2f})_{}_{}'.format(r, train.dataset.idx_to_class[class_id], split_name) metrics[m].append(ap) AP_for_recall[r].append(ap) mAP = np.mean(AP_for_class) metrics['mAP_' + split_name].append(mAP) for r in recalls_mAP: mAP = np.mean(AP_for_recall[r]) metrics['mAP(rec_{:.2f})_{}'.format(r, split_name)].append(mAP) # draw boxes pad = cfg['pad'] im = np.zeros((x.shape[0] + pad * 2, x.shape[1] + pad * 2, x.shape[2])) im[pad:-pad, pad:-pad] = x im = draw_bounding_boxes(im, gt_boxes, color=(1, 0, 0), text_color=(1, 0, 0), pad=pad) im = draw_bounding_boxes(im, pred_boxes, color=(0, 1, 0), text_color=(0, 1, 0), pad=pad) imsave(os.path.join(cfg['out_folder'], 'eval_{}'.format(split_name), 'sample_{:05d}.jpg'.format(im_index)), im) im_index += 1 delta = time.time() - tt0 print('Eval Batch {:04d}/{:04d} on split {}, Time : {:.3f}s'.format(batch + 1, len(loader), split_name, delta)) delta = time.time() - t0 metrics['eval_' + split_name + '_time'] = [delta] print('Eval time of {}: {:.4f}s'.format(split_name, delta)) stats = {} for k in sorted(metrics.keys()): v = np.mean(metrics[k]) print('{}: {:.4}'.format(k, v)) stats[k] = v return stats
def train(*, config='config', resume=False): print('Read config "{}"'.format(config)) cfg = _read_config(config) w_loc = cfg['w_loc'] w_classif = cfg['w_classif'] batch_size = cfg['batch_size'] num_epoch = cfg['num_epoch'] image_size = cfg['image_size'] gamma = cfg['gamma'] mean = cfg['mean'] std = cfg['std'] imbalance_strategy = cfg['imbalance_strategy'] out_folder = cfg['out_folder'] negative_per_positive = cfg['negative_per_positive'] if imbalance_strategy == 'class_weight': pos_weight = cfg['pos_weight'] neg_weight = cfg['neg_weight'] nms_iou_threshold = cfg['nms_iou_threshold'] eval_iou_threshold = cfg['eval_iou_threshold'] log_interval = cfg['log_interval'] nms_score_threshold = cfg['nms_score_threshold'] eval_interval = cfg['eval_interval'] aspect_ratios = cfg['aspect_ratios'] nms_topk = cfg['nms_topk'] debug = cfg['debug'] folders = [ 'train', 'eval_train', 'eval_valid', ] for f in folders: try: os.makedirs(os.path.join(out_folder, f)) except OSError: pass if debug: log_interval = 30 # anchor list for each scale (we have 6 scales) anchor_list = _build_anchor_list(cfg) # dataset for train and valid print('Loading dataset anotations...') (train_dataset, valid_dataset), (train_evaluation, valid_evaluation) = _build_dataset( cfg, anchor_list=anchor_list, ) print('Done loading dataset annotations.') if debug: n = 10 train_dataset = SubSample(train_dataset, nb=n) valid_dataset = SubSample(valid_dataset, nb=n) train_evaluation = SubSample(train_evaluation, nb=n) valid_evaluation = SubSample(valid_evaluation, nb=n) assert train_dataset.class_to_idx == valid_dataset.class_to_idx assert train_dataset.idx_to_class == valid_dataset.idx_to_class clfn = lambda l:l # Dataset loaders for full training and full validation train_loader = DataLoader( train_dataset, shuffle=True, batch_size=batch_size, collate_fn=clfn, num_workers=cfg['num_workers'], ) train_evaluation_loader = DataLoader( train_evaluation, batch_size=batch_size, collate_fn=clfn, num_workers=cfg['num_workers'], ) valid_evaluation_loader = DataLoader( valid_evaluation, batch_size=batch_size, collate_fn=clfn, num_workers=cfg['num_workers'], ) nb_classes = len(train_dataset.class_to_idx) bgid = train_dataset.class_to_idx['background'] class_ids = list(set(train_dataset.class_to_idx.values()) - set([bgid])) class_ids = sorted(class_ids) print('Number of training images : {}'.format(len(train_dataset))) print('Number of valid images : {}'.format(len(valid_dataset))) print('Number of classes : {}'.format(nb_classes)) stats_filename = os.path.join(out_folder, 'stats.csv') train_stats_filename = os.path.join(out_folder, 'train_stats.csv') model_filename = os.path.join(out_folder, 'model.th') optimizer_filename = os.path.join(out_folder, 'optimizer.th') if resume: print('Resuming model: {}'.format(model_filename)) model = torch.load(model_filename) model = model.cuda() if os.path.exists(optimizer_filename): print('Resuming optimizer: {}'.format(optimizer_filename)) optimizer = torch.load(optimizer_filename) else: optimizer_cls = getattr(torch.optim, cfg['optim_algo']) optimizer = optimizer_cls(model.parameters(), lr=0, **cfg['optim_params']) if os.path.exists(stats_filename): stats = pd.read_csv(stats_filename).to_dict(orient='list') first_epoch = max(stats['epoch']) else: stats = defaultdict(list) first_epoch = 1 if os.path.exists(train_stats_filename): train_stats = pd.read_csv(train_stats_filename).to_dict(orient='list') else: train_stats = defaultdict(list) use_discrete_coords = model.use_discrete_coords coords_discretization = model.coords_discretization print('Starting from epoch {}'.format(first_epoch)) else: use_discrete_coords = cfg['use_discrete_coords'] if use_discrete_coords: coords_discretization = torch.linspace( cfg['discrete_coords_min'], cfg['discrete_coords_max'], cfg['discrete_coords_nb'] ) num_coords = 4 * len(coords_discretization) # discretization values for each x,y,w,h else: num_coords = 4 coords_discretization = None if 'init_from' in cfg: print('Init from {}'.format(cfg['init_from'])) model = torch.load(cfg['init_from']) else: model_class = getattr(model_module, cfg['model_name']) kw = cfg.get('model_config', {}) model = model_class( num_anchors=list(map(len, aspect_ratios)), num_classes=nb_classes, num_coords=num_coords, **kw, ) optimizer_cls = getattr(torch.optim, cfg['optim_algo']) optimizer = optimizer_cls(model.parameters(), lr=0, **cfg['optim_params']) model.use_discrete_coords = use_discrete_coords model.num_coords = num_coords model.coords_discretization = coords_discretization model = model.cuda() model.transform = valid_dataset.transform model.nb_classes = nb_classes model.aspect_ratios = aspect_ratios model.anchor_list = anchor_list model.image_size = image_size first_epoch = 1 stats = defaultdict(list) train_stats = defaultdict(list) model.nb_updates = 0 model.avg_loss = 0. model.avg_loc = 0. model.avg_classif = 0 model.background_class_id = train_dataset.background_class_id model.class_to_idx = train_dataset.class_to_idx model.mean = mean model.std = std model.config = cfg print(model) classif_loss_name = cfg.get('classif_loss', 'cross_entropy') model.classif_loss_name = classif_loss_name if classif_loss_name == 'cross_entropy': classif_loss = cross_entropy elif classif_loss_name == 'binary_cross_entropy': classif_loss = binary_cross_entropy_with_logits elif classif_loss_name == 'focal_loss': classif_loss = FocalLoss( gamma=cfg.get('focal_loss_gamma', 2), alpha=cfg.get('focal_loss_alpha', None), ) else: raise ValueError('Unknown classif loss : {}'.format(classif_loss_name)) if imbalance_strategy == 'class_weight': class_weight = torch.zeros(nb_classes) class_weight[0] = neg_weight class_weight[1:] = pos_weight class_weight = class_weight.cuda() for epoch in range(first_epoch, num_epoch): epoch_t0 = time.time() model.train() for batch, samples, in enumerate(train_loader): t0 = time.time() X, (Y, Ypred), (bbox_true, bbox_pred), (class_true, class_pred) = _predict(model, samples) # X is batch of image # Y is groundtruth output # Ypred is predicted output # bbox_true are groundtruth bounding boxes extracted from Y # bbox_pred are predicted bounding boxes extracted from Ypred # class_true are groundtruth classes extracted from Y # class_pred are predicted classes extracted from Ypred m = (class_true != bgid).view(-1) ind = m.nonzero().view(-1) bt = bbox_true bp = bbox_pred ct = class_true cp = class_pred N = max(len(ind), 1.0) # localization loss if use_discrete_coords: l_loc = _discrete_coords_loss(bp[ind], bt[ind], coords_discretization) else: l_loc = smooth_l1_loss(bp[ind], bt[ind], size_average=False) / N # classif loss if imbalance_strategy == 'hard_negative_mining': ind = torch.arange(len(ct)) pos = ind[(ct.data.cpu() > 0)].long().cuda() neg = ind[(ct.data.cpu() == 0)].long().cuda() ct_pos = ct[pos] cp_pos = cp[pos] ct_neg = ct[neg] cp_neg = cp[neg] cp_neg_loss = classif_loss(cp_neg, ct_neg, reduce=False) cp_neg_loss = cp_neg_loss.cuda() vals, indices = cp_neg_loss.sort(descending=True) nb = len(ct_pos) * negative_per_positive cp_neg = cp_neg[indices[0:nb]] ct_neg = ct_neg[indices[0:nb]] l_classif = (classif_loss(cp_pos, ct_pos, size_average=False) + classif_loss(cp_neg, ct_neg, size_average=False)) / N elif imbalance_strategy == 'hard_negative_mining_with_sampling': ind = torch.arange(len(ct)) pos = ind[(ct.data.cpu() > 0)].long().cuda() neg = ind[(ct.data.cpu() == 0)].long().cuda() ct_pos = ct[pos] cp_pos = cp[pos] ct_neg = ct[neg] cp_neg = cp[neg] nb = min(len(ct_pos) * negative_per_positive, len(ct_neg)) cp_neg_loss = classif_loss(cp_neg, ct_neg, reduce=False) proba_sel = torch.nn.Softmax(dim=0)(cp_neg_loss) proba_sel = proba_sel.cuda() indices = torch.multinomial(proba_sel, nb) cp_neg = cp_neg[indices] ct_neg = ct_neg[indices] l_classif = (classif_loss(cp_pos, ct_pos, size_average=False) + classif_loss(cp_neg, ct_neg, size_average=False)) / N elif imbalance_strategy == 'undersampling': ind = torch.arange(len(ct)) pos = ind[(ct.data.cpu() > 0)].long().cuda() neg = ind[(ct.data.cpu() == 0)].long().cuda() ct_pos = ct[pos] cp_pos = cp[pos] ct_neg = ct[neg] cp_neg = cp[neg] nb = len(ct_pos) * negative_per_positive inds = torch.from_numpy(np.random.randint(0, len(ct_neg), nb)) inds = inds.long().cuda() ct_neg = ct_neg[inds] cp_neg = cp_neg[inds] l_classif = (classif_loss(cp_pos, ct_pos, size_average=False) + classif_loss(cp_neg, ct_neg, size_average=False)) / N elif imbalance_strategy == 'class_weight': # TODO make it work if classif_loss is "binary_cross_entropy", it does not # work in that case l_classif = classif_loss(cp, ct, weight=class_weight, size_average=False) / N elif imbalance_strategy == 'nothing': l_classif = classif_loss(cp, ct, size_average=False) / N else: raise ValueError('unknown imbalance strategy : {}'.format(imbalance_strategy)) model.zero_grad() loss = w_loc * l_loc + w_classif * l_classif loss.backward() _update_lr(optimizer, model.nb_updates, cfg['lr_schedule']) optimizer.step() model.avg_loss = model.avg_loss * gamma + item(loss) * (1 - gamma) model.avg_loc = model.avg_loc * gamma + item(l_loc) * (1 - gamma) model.avg_classif = model.avg_classif * gamma + item(l_classif) * (1 - gamma) delta = time.time() - t0 print('Epoch {:05d}/{:05d} Batch {:05d}/{:05d} Loss : {:.3f} Loc : {:.3f} ' 'Classif : {:.3f} AvgTrainLoss : {:.3f} AvgLoc : {:.3f} ' 'AvgClassif {:.3f} Time:{:.3f}s'.format( epoch, num_epoch, batch + 1, len(train_loader), item(loss), item(l_loc), item(l_classif), model.avg_loss, model.avg_loc, model.avg_classif, delta )) train_stats['loss'].append(item(loss)) train_stats['loc'].append(item(l_loc)) train_stats['classif'].append(item(l_classif)) train_stats['time'].append(delta) if model.nb_updates % log_interval == 0: # reporting part # -- draw training samples with their predicted and true bounding boxes pd.DataFrame(train_stats).to_csv(train_stats_filename, index=False) t0 = time.time() torch.save(model, model_filename) torch.save(optimizer, optimizer_filename) X = X.data.cpu().numpy() B = [[b for b, c in y] for y in Y] B = [(np.array(b)) for b in B] C = [[c for b, c in y] for y in Y] C = [(np.array(c)) for c in C] # B contains groundtruth bounding boxes for each scale # C contains groundtruth classes for each scale BP = [ bp.data.cpu().view(bp.size(0), -1, model.num_coords, bp.size(2), bp.size(3)).permute(0, 3, 4, 1, 2).numpy() for bp, cp in Ypred] CP = [ cp.data.cpu().view(cp.size(0), -1, model.nb_classes, cp.size(2), cp.size(3)).permute(0, 3, 4, 1, 2).numpy() for bp, cp in Ypred] # BP contains predicted bounding boxes for each scale # CP contains predicted classes for each scale for i in range(len(X)): # for each example i in mini-batch x = X[i] x = x.transpose((1, 2, 0)) x = x * np.array(std) + np.array(mean) x = x.astype('float32') gt_boxes = [] pred_boxes = [] for j in range(len(Y)): # for each scale j ct = C[j][i]# class_id cp = CP[j][i]#cl1_score,cl2_score,... bt = B[j][i]#4 bp = BP[j][i]#4 if use_discrete_coords: bp = _get_coords(bp, model.coords_discretization) A = model.anchor_list[j] # get groundtruth boxes gt_boxes.extend(decode_bounding_box_list( bt, ct, A, include_scores=False, image_size=image_size, variance=cfg['variance'], )) # get predicted boxes cp = get_probas(cp, classif_loss_name, axis=3) pred_boxes.extend(decode_bounding_box_list( bp, cp, A, include_scores=True, image_size=image_size, variance=cfg['variance'], )) gt_boxes = [(box, class_id) for box, class_id in gt_boxes if class_id != bgid] # apply non-maximal suppression to predicted boxes # 1) for each class, filter low confidence predictions and then do NMS # 2) concat all the bboxes from all classes # 3) take to the topk (nms_topk) if cfg['use_nms']: pred_boxes = non_maximal_suppression_per_class( pred_boxes, background_class_id=bgid, iou_threshold=nms_iou_threshold, score_threshold=nms_score_threshold) pred_boxes = pred_boxes[0:nms_topk] else: pred_boxes = [ (box, scores.argmax(), scores.max()) for box, scores in pred_boxes if scores.argmax() != bgid and scores.max() > nms_score_threshold ] # get class names gt_boxes = [(box, train_dataset.idx_to_class[class_id]) for box, class_id in gt_boxes] pred_boxes = [(box, train_dataset.idx_to_class[class_id], score) for box, class_id, score in pred_boxes] # draw boxes pad = cfg['pad'] im = np.zeros((x.shape[0] + pad * 2, x.shape[1] + pad * 2, x.shape[2])) im[pad:-pad, pad:-pad] = x im = draw_bounding_boxes(im, gt_boxes, color=(1, 0, 0), text_color=(1, 0, 0), pad=pad) im = draw_bounding_boxes(im, pred_boxes, color=(0, 1, 0), text_color=(0, 1, 0), pad=pad) imsave(os.path.join(out_folder, 'train', 'sample_{:05d}.jpg'.format(i)), im) delta = time.time() - t0 print('Draw box time {:.4f}s'.format(delta)) model.nb_updates += 1 epoch_time = time.time() - epoch_t0 if epoch % eval_interval != 0: continue if cfg.get('evaluate', True) is False: continue stats_epoch = _evaluate_model(model, train_evaluation_loader, valid_evaluation_loader) stats_epoch['train_time'] = epoch_time stats_epoch['epoch'] = epoch for k, v in stats_epoch.items(): stats[k].append(v) pd.DataFrame(stats).to_csv(stats_filename, index=False)
import bounding_box import cv2 as cv import matplotlib.pyplot as plt import numpy as np import sys image = sys.argv[1] + '.jpg' image = cv.imread(image) image = np.asarray(cv.cvtColor(image, cv.COLOR_BGR2GRAY)) bounding_box.plot(image) boxes = bounding_box.get_bounding_boxes(image) print(boxes) bounding_box.draw_bounding_boxes(image, boxes) bounding_box.plot(image)