def main(): print('Loading model from %s' % (args.ckpt)) net = RetinaNet() load_checkpoint('%s' % (args.ckpt), net) net.eval() use_cuda = torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") print('Device used:', device) net = net.to(device) encoder = DataEncoder() if args.vmode: if not os.path.exists("%s" % (args.pred_root)): os.makedirs("%s" % (args.pred_root)) visualize(encoder, net, device, args.image_root, val_image_list, args.pred_root) if args.mmode: if not os.path.exists("./mPA"): os.makedirs("./mPA") if os.path.exists("./mPA/detection-results"): print("Remove Pred file") os.system("rm -rf ./mPA/detection-results") os.makedirs("./mPA/detection-results") if not os.path.exists("./mPA/detection-results"): os.makedirs("./mPA/detection-results") mPA_pred(encoder, net, device, args.image_root, args.anno_root) return 0
def make_model(): resnet50_backbone = get_backbone() loss_fn = RetinaNetLoss(num_classes) model = RetinaNet(num_classes, resnet50_backbone) optimizer = tf.optimizers.SGD(learning_rate=learning_rate_fn, momentum=0.9) model.compile(loss=loss_fn, optimizer=optimizer) return model
def run_train(args): assert torch.cuda.is_available(), 'Error: CUDA not found!' best_loss = float('inf') # best test loss start_epoch = 0 # start from epoch 0 or last epoch # Data print('==> Preparing data..') trainloader = get_train_loader(img_dir=settings.IMG_DIR, batch_size=batch_size) #trainloader = get_small_train_loader() print(trainloader.num) #testloader = get_train_loader(img_dir=settings.IMG_DIR) # Model net = RetinaNet() #net.load_state_dict(torch.load('./model/net.pth')) net.load_state_dict(torch.load('./ckps/best_0.pth')) net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() criterion = FocalLoss() #optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) optimizer = optim.Adam(net.parameters(), lr=args.lr) iter_save = 200 bgtime = time.time() # Training for epoch in range(start_epoch, start_epoch + 100): print('\nEpoch: %d' % epoch) net.train() #net.module.freeze_bn() train_loss = 0 for batch_idx, (inputs, loc_targets, cls_targets) in enumerate(trainloader): inputs = Variable(inputs.cuda()) loc_targets = Variable(loc_targets.cuda()) cls_targets = Variable(cls_targets.cuda()) optimizer.zero_grad() loc_preds, cls_preds = net(inputs) loss = criterion(loc_preds, loc_targets, cls_preds, cls_targets) loss.backward() optimizer.step() #train_loss += loss.data[0] sample_num = (batch_idx + 1) * batch_size avg_loss = running_loss(loss.data[0]) print( 'Epoch: {}, num: {}/{} train_loss: {:.3f} | run_loss: {:.3f} min: {:.1f}' .format(epoch, sample_num, trainloader.num, loss.data[0], avg_loss, (time.time() - bgtime) / 60), end='\r') if batch_idx % iter_save == 0: torch.save( net.module.state_dict(), './ckps/best_{}.pth'.format(batch_idx // iter_save % 5)) log.info('batch: {}, loss: {:.4f}'.format(batch_idx, avg_loss))
def __init__(self, device='cuda', verbose=False): self.verbose = verbose self.net = RetinaNet(backbone=cfg.backbone, num_classes=2, pretrained=False) checkpoint = torch.load(os.path.join('ckpts', 'retina_fp16_2e4', '12_ckpt.pth'), map_location=device) errors = self.net.load_state_dict(checkpoint['net']) logging.warning('Errors from loading Retina model: {}'.format(errors)) self.net = self.net.half().eval().to(device) self.net.device = device
def train(total_epochs=1, interval=100, resume=False, ckpt_path = ''): print("Loading training dataset...") train_dset = OpenImagesDataset(root='./data/train', list_file ='./data/tmp/train_images_bbox.csv', transform=transform, train=True, input_size=600) train_loader = data.DataLoader(train_dset, batch_size=4, shuffle=True, num_workers=4, collate_fn=train_dset.collate_fn) print("Loading completed.") #val_dset = OpenImagesDataset(root='./data/train', # list_file='./data/tmp/train_images_bbox.csv', train=False, transform=transform, input_size=600) #val_loader = torch.utils.data.DataLoader(val_dset, batch_size=1, shuffle=False, num_workers=4, collate_fn=val_dset.collate_fn) net = RetinaNet() net.load_state_dict(torch.load('./model/net.pth')) criterion = FocalLoss() net.cuda() criterion.cuda() optimizer = optim.SGD(net.parameters(), lr=1e-3, momentum=0.9, weight_decay=1e-4) best_val_loss = 1000 start_epoch=0 if resume: if os.path.isfile(ckpt_path): print(f'Loading from the checkpoint {ckpt_path}') checkpoint = torch.load(ckpt_path) start_epoch = checkpoint['epoch'] best_val_loss = checkpoint['best_val_loss'] net.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print(f'Loaded checkpoint {ckpt_path}, epoch : {start_epoch}') else: print(f'No check point found at the path {ckpt_path}') for epoch in range(start_epoch, total_epochs): train_one_epoch(train_loader, net, criterion, optimizer, epoch, interval) val_loss = 0 #val_loss = validate(val_loader, net, criterion, interval) if val_loss < best_val_loss: best_val_loss = val_loss save_checkpoint({ 'epoch': epoch+1, 'state_dict': net.state_dict(), 'best_val_loss': best_val_loss, 'optimizer' : optimizer.state_dict() }, is_best=True)
def test_retinanet(): """test retinanet""" net = RetinaNet(classes=80) x = Variable(torch.rand(1, 3, 500, 500), volatile=True) now = time.time() net.cuda() predictions = net(x) later = time.time() print(later - now) for prediction in predictions: print(prediction.size())
def main(_): assert tfe.num_gpus() > 0, 'Make sure the GPU device exists' device_name = '/gpu:{}'.format(args.cuda_device) print('\n==> ==> ==> Using device {}'.format(device_name)) # Load the dataset train_ds, val_ds = [ dataset_generator( mode, conf.input_size, num_epochs=1, batch_size=conf.batch_size, buffer_size=10000) # TODO edit this when in real training for mode in ['train', 'val'] ] # Create the model and optimizer model = RetinaNet() optimizer = tf.train.RMSPropOptimizer(conf.learning_rate) # Define the path to the TensorBoard summary train_dir, val_dir = [ os.path.join(conf.summary_dir, mode) for mode in ['train', 'val'] ] tf.gfile.MakeDirs(conf.summary_dir) train_summary_writer = tf.contrib.summary.create_summary_file_writer( train_dir, flush_millis=10000, name='train') val_summary_writer = tf.contrib.summary.create_summary_file_writer( val_dir, flush_millis=10000, name='val') checkpoint_prefix = os.path.join(conf.checkpoint_dir, 'ckpt') with tfe.restore_variables_on_create( tf.train.latest_checkpoint(conf.checkpoint_dir)): with tf.device(device_name): epoch = tfe.Variable(1., name='epoch') best_loss = tfe.Variable(tf.float32.max, name='best_loss') print('==> ==> ==> Start training from epoch {:.0f}...\n'.format( epoch.numpy())) while epoch <= conf.num_epochs + 1: gs = tf.train.get_or_create_global_step() with train_summary_writer.as_default(): train_one_epoch(model, optimizer, train_ds, epoch.numpy()) with val_summary_writer.as_default(): eval_loss = validate(model, val_ds, epoch.numpy()) # Save the best loss if eval_loss < best_loss: best_loss.assign( eval_loss) # do NOT be copied directly, SHALLOW! all_variables = (model.variables + optimizer.variables() + [gs] + [epoch] + [best_loss]) tfe.Saver(all_variables).save(checkpoint_prefix, global_step=gs) epoch.assign_add(1)
def main(): print('==> chooseing data..') transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # Model net = RetinaNet() criterion = FocalLoss() # optimizer = optim.Adam(net.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-4) optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9, weight_decay=1e-4) load_model_epoch = args.load_model_epoch checkpoint = torch.load( './checkpoint/{}_ckpt.pth'.format(load_model_epoch)) # max_epoch net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() net.load_state_dict(checkpoint['net']) optimizer.load_state_dict(checkpoint['optimizer']) start_epoch = checkpoint['epoch'] test(start_epoch, transform, net, criterion, optimizer)
def main(): # assert torch.cuda.is_available(), 'Error: CUDA not found!' best_loss = float('inf') # best test loss start_epoch = 0 # start from epoch 0 or last epoch save_model_path = args.model # Data print('==> Preparing data..') transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # Model net = RetinaNet() net.load_state_dict(torch.load('./model/net.pth')) net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() criterion = FocalLoss() # optimizer = optim.Adam(net.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-4) optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) for epoch in range(start_epoch, start_epoch + args.train_epoch): train(epoch + 1, transform, net, optimizer, criterion) save_model(epoch + 1, save_model_path, net, optimizer)
def build_model(cfg): cfg.build_backbone = build_backbone cfg.build_anchor_generator = build_anchor_generator model = RetinaNet(cfg) logger = logging.getLogger(__name__) logger.info("Model:\n{}".format(model)) return model
class RetinaDetector: def __init__(self, device='cuda', verbose=False): self.verbose = verbose self.net = RetinaNet(backbone=cfg.backbone, num_classes=2, pretrained=False) checkpoint = torch.load(os.path.join('ckpts', 'retina_fp16_2e4', '12_ckpt.pth'), map_location=device) errors = self.net.load_state_dict(checkpoint['net']) logging.warning('Errors from loading Retina model: {}'.format(errors)) self.net = self.net.half().eval().to(device) self.net.device = device def detect(self, img): with torch.no_grad(): boxes, labels, scores = hires_prediction(img, self.net, verbose=self.verbose) return boxes, labels, scores
def train_obj(): # Model net = RetinaNet() net = torch.nn.DataParallel(net, device_ids=range( torch.cuda.device_count())) net.cuda() criterion = FocalLoss() optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) scheduler_obj = torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambda=lambda x: (1 - x / (len(trainloader) * epochs))**0.9) obj_trainer = ObjDetTrainer(net, criterion, optimizer, scheduler_obj, trainloader, valloader, device) obj_trainer.train(epochs, True)
def _create_losses(input_queue, num_classes, train_config): """Creates loss function for a DetectionModel. Args: input_queue: BatchQueue object holding enqueued tensor_dicts. num_classes: num of classes, integer Returns: Average sum of loss of given input batch samples with shape """ (images, groundtruth_boxes_list, groundtruth_classes_list, anchors_list) = _get_inputs(input_queue, num_classes, batch_size=train_config.batch_size) images = [ preprocess(image, im_height=train_config.im_height, im_width=train_config.im_width, preprocess_options=train_config.data_augmentation_ops) for image in images ] images = tf.concat(images, 0) net = RetinaNet() loc_preds, cls_preds = net(images, num_classes + 1, anchors=9) # get num of anchor overlapped with ground truth box cls_gt = [anchor.get_field("gt_labels") for anchor in anchors_list] loc_gt = [anchor.get_field("gt_encoded_boxes") for anchor in anchors_list] # pos anchor count for each image gt_anchor_nums = tf.map_fn( lambda x: tf.reduce_sum(tf.cast(tf.greater(x, 0), tf.int32)), cls_gt) # get valid anchor indices valid_anchor_indices = tf.squeeze(tf.where(tf.greater_equal(cls_gt, 0))) # skip ignored anchors (iou belong to 0.4 to 0.5) [valid_cls_preds, valid_cls_gt] = map(lambda x: tf.gather(x, valid_anchor_indices, axis=1), [cls_preds, cls_gt]) # classification loss: convert to onehot code cls_loss = tf.multiply(focal_loss(valid_cls_gt, valid_cls_preds), 1. / tf.to_float(gt_anchor_nums)) # location regression loss valid_cls_indices = tf.squeeze(tf.where(tf.greater(cls_gt, 0))) # skip negative and ignored anchors [valid_loc_preds, valid_loc_gt] = map(lambda x: tf.gather(x, valid_cls_indices, axis=1), [loc_preds, loc_gt]) loc_loss = regression_loss(valid_loc_preds, valid_loc_gt, weights=tf.expand_dims( 1. / tf.to_float(gt_anchor_nums), 1)) loss = (tf.reduce_sum(loc_loss) + tf.reduce_sum(cls_loss)) / tf.size( gt_anchor_nums, out_type=tf.float32) return loss
def load_model(backbone): print('loading model...') model = torch.load(os.path.join('model', 'restnet101_8K.pth')) net = RetinaNet(backbone=backbone, num_classes=len(cfg.classes)) net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() cudnn.benchmark = True net.load_state_dict(model['net']) return net
def evaluate_threshold(img_ids, cls_threshold, bbox_dict): dloader = get_test_loader(img_ids, img_dir=settings.IMG_DIR, batch_size=batch_size) # Model net = RetinaNet() net.load_state_dict(torch.load(CKP_FILE)) #net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() net.eval() bgtime = time.time() encoder = DataEncoder() encoder.class_threshold = cls_threshold true_objects_num = 0 pred_objects_num = 0 for batch_idx, inputs in enumerate(dloader): inputs = Variable(inputs.cuda()) loc_preds, cls_preds = net(inputs) for i in range(len(loc_preds)): boxes, labels, scores = encoder.decode( loc_preds[i].data, cls_preds[i].data, (settings.IMG_SZ, settings.IMG_SZ)) pred_objects_num += len(scores) for img_idx in range(len(inputs)): img_id = dloader.img_ids[batch_idx * batch_size + img_idx] if img_id in bbox_dict: true_objects_num += len(bbox_dict[img_id]) print('{} / {}, {} / {}, {:.4f}, {:.2f} min'.format( batch_size * (batch_idx + 1), dloader.num, pred_objects_num, true_objects_num, cls_threshold, (time.time() - bgtime) / 60), end='\r') print('\n') print('=>>> {}/{}, {}, {:.4f}\n'.format( pred_objects_num, true_objects_num, pred_objects_num - true_objects_num, cls_threshold))
def predict(): assert torch.cuda.is_available(), 'Error: CUDA not found!' print('==> Preparing data..') dloader = get_test_loader(get_test_ids(), img_dir=settings.TEST_IMG_DIR, batch_size=batch_size) print(dloader.num) # Model net = RetinaNet() net.load_state_dict(torch.load(CKP_FILE)) #net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() bgtime = time.time() encoder = DataEncoder() net.eval() prediction_strings = [] for batch_idx, inputs in enumerate(dloader): inputs = Variable(inputs.cuda()) loc_preds, cls_preds = net(inputs) print('{} / {} {:.2f}'.format(batch_size * (batch_idx + 1), dloader.num, (time.time() - bgtime) / 60), end='\r') for i in range(len(loc_preds)): boxes, labels, scores = encoder.decode( loc_preds[i].data, cls_preds[i].data, (settings.IMG_SZ, settings.IMG_SZ)) prediction_strings.append( _get_prediction_string(boxes, labels, scores)) print(len(prediction_strings)) print(prediction_strings[:3]) submission = pd.DataFrame({ 'ImageId': dloader.img_ids, 'PredictionString': prediction_strings }) submission.to_csv('sub7.csv', index=False)
def test(): input_size = (448, 448) dataset = dataset_generator('val', input_size, 1, 16, 100) box_encoder = BoxEncoder() model = RetinaNet() with tf.device("gpu:0"): for i, (image, loc_trues, cls_trues) in enumerate(tfe.Iterator(dataset)): print('loc_trues shape: {}'.format(loc_trues.shape)) print('cls_trues shape: {}'.format(cls_trues.shape)) loc_preds, cls_preds = model(image) with tf.device("cpu:0"): # Decode one by one in a batch loc_preds, cls_preds, score = box_encoder.decode_batch( loc_preds.cpu(), cls_preds.cpu(), input_size) print('loc_preds {} shape: {}'.format(loc_preds, loc_preds.shape)) print('cls_preds {} shape: {}'.format(cls_preds, cls_preds.shape)) print('score {} shape: {}'.format(score, score.shape)) break
def test(): # with tf.device("/gpu:0"): # [batch_size, #anchors]s # loc_preds = tf.random_uniform([3, 10, 4]) # loc_trues = tf.random_uniform([3, 10, 4]) # cls_preds = tf.random_uniform([3, 10, 12]) # cls_trues = tf.random_uniform([3, 10]) from inputs import dataset_generator from retinanet import RetinaNet dataset = dataset_generator('train', (448, 448), 1, 16, 100) model = RetinaNet() with tf.device("/gpu:0"): for i, (image, loc_trues, cls_trues) in enumerate(tfe.Iterator(dataset)): loc_preds, cls_preds = model(image, training=True) loc_loss, cls_loss = loss_fn(loc_preds, loc_trues, cls_preds, cls_trues) print( "Step 0: Location loss: {:.5f} | Class loss: {:.5f}".format( loc_loss.numpy(), cls_loss.numpy())) break
num_workers=8, collate_fn=trainset.collate_fn) testset = ListDataset( list_file='/group/proteindl/ps793/Dota/val_retina_dota_nd.txt', train=False, transform=transform, input_size=1024) testloader = torch.utils.data.DataLoader(testset, batch_size=2, shuffle=False, num_workers=8, collate_fn=testset.collate_fn) # Model net = RetinaNet(num_classes=15) net.load_state_dict(torch.load('./model/dota_15c_9ma.pth')) if resume: print('==> Resuming from checkpoint..') checkpoint = torch.load('./checkpoint/ckpt.pth') net.load_state_dict(checkpoint['net']) best_loss = checkpoint['loss'] start_epoch = checkpoint['epoch'] net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() if fix == 'head': for param in net.module.fpn.conv1.parameters(): param.requires_grad = False for param in net.module.fpn.bn1.parameters():
transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) trainset = NRF(root='./data', image_set=args.type, transform=transform, input_size=768) trainloader = torch.utils.data.DataLoader(trainset, batch_size=16, shuffle=True, num_workers=4, collate_fn=trainset.collate_fn) ckpt_dir = os.path.join('./checkpoint_list_wise/albertsons', args.ckpt) if not os.path.exists(ckpt_dir): os.makedirs(ckpt_dir) # Model print('==> Setting up network..') if args.type == 'prod': print('im gonno use Retina_Net') net = RetinaNet(num_classes=1, num_anchors=15, backbone=args.net) elif args.type == 'tag': net = TagNet(num_classes=1, num_anchors=9) else: raise TypeError('Unknown detection type, choose "prod" or "tag"') if args.resume: print('==> Resuming from checkpoint..') checkpoint = torch.load('/media/Darius/shayeree/training/checkpoint_list_wise/albertsons/Retina50ProdB1/ckpt_0005_0.1237.pth') net.load_state_dict(checkpoint['net']) best_loss = checkpoint['loss'] start_epoch = checkpoint['epoch'] else: params_dict = torch.load('./model/{:s}.pth'.format(args.net)) net_dict = net.fpn.state_dict() params_dict = {k: v for k, v in params_dict.items() if k in net_dict}
# model backbone if args.backbone == 'resnet18': backbone = resnet18_features elif args.backbone == 'resnet34': backbone = resnet34_features elif args.backbone == 'resnet50': backbone = resnet50_features elif args.backbone == 'resnet101': backbone = resnet101_features elif args.backbone == 'resnet152': backbone = resnet152_features else: raise ValueError('Invalid model backbone') # RetinaNet model model = RetinaNet(80, backbone, pretrained=True) optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0001) # TODO: enable more flexible lr profile scheduler = lr_scheduler.MultiStepLR( optimizer, # Milestones are set assuming batch size is 16: # 60000 / batch_size = 3750 # 80000 / batch_size = 5000 milestones=[3750, 5000], gamma=0.1)
print('==> Preparing data..') transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485,0.456,0.406), (0.229,0.224,0.225)) ]) trainset = ListDataset(root='/search/odin/liukuang/data/voc_all_images', list_file='./data/voc12_train.txt', train=True, transform=transform, input_size=600) trainloader = torch.utils.data.DataLoader(trainset, batch_size=16, shuffle=True, num_workers=8, collate_fn=trainset.collate_fn) testset = ListDataset(root='/search/odin/liukuang/data/voc_all_images', list_file='./data/voc12_val.txt', train=False, transform=transform, input_size=600) testloader = torch.utils.data.DataLoader(testset, batch_size=16, shuffle=False, num_workers=8, collate_fn=testset.collate_fn) # Model net = RetinaNet() net.load_state_dict(torch.load('./model/net.pth')) if args.resume: print('==> Resuming from checkpoint..') checkpoint = torch.load('./checkpoint/ckpt.pth') net.load_state_dict(checkpoint['net']) best_loss = checkpoint['loss'] start_epoch = checkpoint['epoch'] net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() criterion = FocalLoss() optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) # Training
from fpn import FPN50 from retinanet import RetinaNet print('Loading pretrained ResNet50 model..') d = torch.load('./model/resnet50.pth') print('Loading into FPN50..') fpn = FPN50() dd = fpn.state_dict() for k in d.keys(): if not k.startswith('fc'): # skip fc layers dd[k] = d[k] print('Saving RetinaNet..') net = RetinaNet() for m in net.modules(): if isinstance(m, nn.Conv2d): init.normal(m.weight, mean=0, std=0.01) if m.bias is not None: init.constant(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): m.weight.data.fill_(1) m.bias.data.zero_() pi = 0.01 init.constant(net.cls_head[-1].bias, -math.log((1 - pi) / pi)) net.fpn.load_state_dict(dd) torch.save(net.state_dict(), 'net.pth') print('Done!')
#-------------------------------------# # 调用摄像头检测 #-------------------------------------# from retinanet import RetinaNet from PIL import Image import numpy as np import cv2 import time retinanet = RetinaNet() # 调用摄像头 capture = cv2.VideoCapture(0) # capture=cv2.VideoCapture("1.mp4") fps = 0.0 while (True): t1 = time.time() # 读取某一帧 ref, frame = capture.read() # 格式转变,BGRtoRGB frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # 转变成Image frame = Image.fromarray(np.uint8(frame)) # 进行检测 frame = np.array(efficientdet.detect_image(frame)) # RGBtoBGR满足opencv显示格式 frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) fps = (fps + (1. / (time.time() - t1))) / 2 print("fps= %.2f" % (fps))
''' predict.py有几个注意点 1、无法进行批量预测,如果想要批量预测,可以利用os.listdir()遍历文件夹,利用Image.open打开图片文件进行预测。 2、如果想要保存,利用r_image.save("img.jpg")即可保存。 3、如果想要获得框的坐标,可以进入detect_image函数,读取top,left,bottom,right这四个值。 4、如果想要截取下目标,可以利用获取到的top,left,bottom,right这四个值在原图上利用矩阵的方式进行截取。 ''' from retinanet import RetinaNet from PIL import Image retinanet = RetinaNet() while True: img = input('Input image filename:') try: image = Image.open(img) except: print('Open Error! Try again!') continue else: r_image = retinanet.detect_image(image) r_image.show()
print('==> Preparing data..') transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485,0.456,0.406), (0.229,0.224,0.225)) ]) trainset = ListDataset(root='/search/odin/liukuang/data/voc_all_images', list_file='./voc_data/test.txt', train=True, transform=transform, input_size=600) trainloader = torch.utils.data.DataLoader(trainset, batch_size=16, shuffle=True, num_workers=8, collate_fn=trainset.collate_fn) testset = ListDataset(root='/search/odin/liukuang/data/voc_all_images', list_file='./voc_data/test.txt', train=False, transform=transform, input_size=600) testloader = torch.utils.data.DataLoader(testset, batch_size=16, shuffle=False, num_workers=8, collate_fn=testset.collate_fn) # Model net = RetinaNet() net.load_state_dict(torch.load('./model/net.pth')) if args.resume: print('==> Resuming from checkpoint..') checkpoint = torch.load('./checkpoint/ckpt.pth') net.load_state_dict(checkpoint['net']) best_loss = checkpoint['loss'] start_epoch = checkpoint['epoch'] net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() criterion = FocalLoss() optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) # Training
collate_fn=trainset.collate_fn) testset = ListDataset( root='/home/deva/s165/RFBNet-master/data/VOCdevkit/VOC2008/JPEGImages', list_file='./data/test.txt', train=False, transform=transform, input_size=512) testloader = torch.utils.data.DataLoader(testset, batch_size=8, shuffle=False, num_workers=8, collate_fn=testset.collate_fn) # Model net = RetinaNet() net.load_state_dict(torch.load('./model/net.pth')) args.resume = True if args.resume: print('==> Resuming from checkpoint..') checkpoint = torch.load('./checkpoint/ckpt.pth') net.load_state_dict(checkpoint['net']) # best_loss = checkpoint['loss'] start_epoch = checkpoint['epoch'] net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() criterion = FocalLoss() optimizer = optim.SGD(net.parameters(), lr=args.lr,
from fpn import FPN50 from retinanet import RetinaNet print('Loading pretrained ResNet50 model..') d = torch.load('./model/resnet50.pth') print('Loading into FPN50..') fpn = FPN50() dd = fpn.state_dict() for k in d.keys(): if not k.startswith('fc'): # skip fc layers dd[k] = d[k] print('Saving RetinaNet..') net = RetinaNet() for m in net.modules(): if isinstance(m, nn.Conv2d): init.normal(m.weight, mean=0, std=0.01) if m.bias is not None: init.constant(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): m.weight.data.fill_(1) m.bias.data.zero_() pi = 0.01 init.constant(net.cls_head[-1].bias, -math.log((1-pi)/pi)) net.fpn.load_state_dict(dd) torch.save(net.state_dict(), 'net.pth') print('Done!')
import torchvision.utils as utils from torch.autograd import Variable from fcn_inception_smallMap import FCN from retinanet import RetinaNet from encoder import DataEncoder from PIL import Image, ImageDraw import os import cv2 import time import math import numpy as np os.environ["CUDA_VISIBLE_DEVICES"] = "1" print('Loading model..') net_det = RetinaNet().cuda() net_det.load_state_dict( torch.load('../pytorch-retinanet/checkpoint/model_10_.pth')['net']) net_det.eval() net_det = torch.nn.DataParallel(net_det, device_ids=range(torch.cuda.device_count())) net_map = FCN().cuda() net_map.load_state_dict( torch.load('./../dense_finetune/models/inception_smallMap_3_0.8880.pth')) net_map.eval() net_map = torch.nn.DataParallel(net_map, device_ids=range(torch.cuda.device_count())) transform = transforms.Compose([ transforms.ToTensor(),
fontScale = 0.4 thickness = 1 t_size, _ = cv2.getTextSize(text, fontFace, fontScale + 0.1, thickness) p_org = tuple(np.amax(poly['poly'], 0)) p_top = (p_org[0] - t_size[0], p_org[1] - t_size[1]) p = (p_top[0] + 1, p_org[1] - 1) cv2.rectangle(img, p_top, p_org, bg, cv2.FILLED) cv2.putText(img, text, p, fontFace, fontScale, [255, 255, 255], 1, cv2.LINE_AA) return img batch_size = 16 network = RetinaNet("eval_thn", project_dir="/root/retinanet").cuda() network.load_state_dict( torch.load( "/root/retinanet/training_logs/model_13/checkpoints/model_13_epoch_50.pth" )) test_dataset = DatasetThnSeqSynscapes(thn_data_path="/root/deeplabv3/data/thn") #test_dataset = DatasetThnSeq(thn_data_path="/root/deeplabv3/data/thn") bbox_encoder = BboxEncoder(img_h=test_dataset.img_height, img_w=test_dataset.img_width) num_test_batches = int(len(test_dataset) / batch_size) print("num_test_batches:", num_test_batches) test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
import torch import torchvision.transforms as transforms from torchsummary import summary from torch.autograd import Variable from retinanet import RetinaNet from encoder import DataEncoder from PIL import Image, ImageDraw import cv2 import time print('Loading model..') net = RetinaNet() net.load_state_dict(torch.load('checkpoint/ckpt.pth')['net']) net.eval() net.cuda() #summary(net, (3, 224, 224)) transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) print('Loading image..') #img = Image.open('./image/000001.jpg') cap = cv2.VideoCapture('Z:\\u_zhh\\video\\jy2.mp4') while True: ret, cv2img = cap.read() # 名称不能有汉字
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_sizes, shuffle=False, num_workers=8, collate_fn=testset.collate_fn) encoder = DataEncoder() evaluation(model, batch_sizes, testloader, the_classes, encoder, thresh, im_size, im_size) map = eval(the_classes) return map if __name__ == '__main__': batch_sizes = 1 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) test_path = '/home/ecust/gmy/pytorch-retinanet-master/pytorch-retinanet-master/data/NEU-DET/valid/labels' print('Loading model..') net = RetinaNet() # net = torch.nn.DataParallel(net) # cudnn.benchmark = True net.load_state_dict(torch.load('checkpoints/ckpt_208.pth')) # net.load_state_dict(torch.load('weights/ckpt_1175.pth')) evaluate(net, transform, test_path, batch_sizes, thresh=0.001, im_size=300)