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)
