import os import argparse import os import cv2 import torch import argparse from datasets import build_dataset, get_coco_api_from_dataset from models import build_test_model from datasets.coco_eval import CocoEvaluator import numpy as np import util.misc as utils def test_img(args, model, postprocessors, save_path): dataset_test = build_dataset( image_set="val", args=args, eval_in_training_set=False, ) sampler_test = torch.utils.data.SequentialSampler(dataset_test) data_loader_test = torch.utils.data.DataLoader( dataset_test, 1, sampler=sampler_test, drop_last=False, collate_fn=utils.collate_fn, num_workers=0, pin_memory=True,) base_ds = get_coco_api_from_dataset(dataset_test) for img_data, targets in data_loader_test: img_data = img_data.to(device) targets = [{k: v.to(device) for k, v in t.items()} for t in targets] outputs = model(img_data) orig_target_sizes = torch.stack([t["orig_size"] for t in target], dim=0) result = postprocessors['bbox'](outputs, orig_target_sizes) if 'segm' in postprocessors.keys(): target_sizes = torch.stack([t["size"] for t in targets], dim=0) results = postprocessors['segm'](results, outputs, orig_target_sizes, target_sizes) res = {target['image_id'].item(): output for target, output in zip(targets, result)} iou_types = tuple(k for k in ("segm", "bbox") if k in postprocessors.keys()) coco_evaluator = CocoEvaluator(base_ds, iou_types) if coco_evaluator is not None: coco_evaluator.update(res) res = res[targets[0]['image_id'].item()] min_score = 0.65 img_name = dataset_test.coco.loadImgs(targets[0]['image_id'].item())[0]['file_name'] img = cv2.imread(os.path.join(args.coco_path, 'images/val2017', img_name)) draw_img = img.copy() save_status = False for i in range(0, 100): res_tmp = res['scores'] if float(res_tmp[i]) > min_score: save_status = True score = float(res_tmp[i]) label = int(res['labels'][i].cpu().numpy()) bbox = res['boxes'][i].cpu().numpy().tolist() print("***", label, bbox) cv2.putText(draw_img, "{} | {}".format(label, str(score)[:3]), (int(bbox[0]), int(bbox[1])-2), 0, 0.5, (255, 255, 255), 1) cv2.rectangle(draw_img, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (0, 0, 255), 1) if save_status: cv2.imwrite("{}/{}".format(save_path, img_name), draw_img) if coco_evaluator is not None: coco_evaluator.synchronize_between_processes() coco_evaluator.accumulate() coco_evaluator.summarize() print(coco_evaluator) def get_parser(): parser = argparse.ArgumentParser("DETR Detector", add_help=False) parser.add_argument('--dataset_file', default='coco') parser.add_argument('--coco_path', default='../datasets/COCO2017', type=str) parser.add_argument('--coco_panoptic_path', type=str) parser.add_argument('--remove_difficult', action='store_true') parser.add_argument('--output_dir', default='', help='path where to save, empty for no saving') parser.add_argument('--device', default='cuda:0', help='device to use for training / testing') parser.add_argument('--seed', default=42, type=int) # * Segmentation parser.add_argument('--masks', action='store_true', help="Train segmentation head if the flag is provided") # * Backbone parser.add_argument('--backbone', default='resnet50', type=str, help="Name of the convolutional backbone to use") parser.add_argument('--dilation', action='store_true', help="If true, we replace stride with dilation in the last convolutional block (DC5)") parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'), help="Type of positional embedding to use on top of the image features") # * Transformer parser.add_argument('--enc_layers', default=6, type=int, help="Number of encoding layers in the transformer") parser.add_argument('--dec_layers', default=6, type=int, help="Number of decoding layers in the transformer") parser.add_argument('--dim_feedforward', default=2048, type=int, help="Intermediate size of the feedforward layers in the transformer blocks") parser.add_argument('--hidden_dim', default=256, type=int, help="Size of the embeddings (dimension of the transformer)") parser.add_argument('--dropout', default=0.1, type=float, help="Dropout applied in the transformer") parser.add_argument('--nheads', default=8, type=int, help="Number of attention heads inside the transformer's attentions") parser.add_argument('--num_queries', default=100, type=int, help="Number of query slots") parser.add_argument('--pre_norm', action='store_true') # Loss parser.add_argument('--no_aux_loss', dest='aux_loss', action='store_false', help="Disables auxiliary decoding losses (loss at each layer)") # * eval technologies parser.add_argument("--eval", action="store_true") # eval in training set parser.add_argument("--eval_in_training_set", default=False, action="store_true") # topk for eval parser.add_argument("--topk", default=100, type=int) # * training technologies parser.add_argument("--use_fp16", default=False, action="store_true") parser.add_argument("--use_checkpoint", default=False, action="store_true") return parser if __name__ == "__main__": args = get_parser().parse_args() device = torch.device(args.device) load_path = './checkpoints/r50_hybrid_branch_lambda1_group6_t1500_dp0_mqs_lft_deformable_detr_plus_iterative_bbox_refinement_plus_plus_two_stage_36eps.pth' img_path = '/home/datasets/COCO2017' save_path = './result' if not os.path.exists(save_path): os.makedirs(save_path) model, postprocessors = build_test_model(args) model.to(device) checkpoint = torch.load(load_path, map_location='cpu') model.load_state_dict(checkpoint["model"], False) model.eval() test_img(args, model, postprocessors, save_path)