action='store_true',
help='resume from checkpoint')
parser.add_argument('--model',
default='./examples/ssd/model/ssd512_vgg16.pth',
type=str,
help='initialized model path')
parser.add_argument('--checkpoint',
default='./examples/ssd/checkpoint/ckpt.pth',
type=str,
help='checkpoint path')
args = parser.parse_args()
# Model
print('==> Building model..')
# net = SSD512(num_classes=21)
net = FPNSSD512(num_classes=81)
# net.load_state_dict(torch.load(args.model))
best_loss = float('inf') # best test loss
start_epoch = 0 # start from epoch 0 or last epoch
if args.resume:
print('==> Resuming from checkpoint..')
checkpoint = torch.load(args.checkpoint)
net.load_state_dict(checkpoint['net'])
best_loss = checkpoint['loss']
start_epoch = checkpoint['epoch']
# Dataset
print('==> Preparing dataset..')
box_coder = SSDBoxCoder(net)
img_size = 512
list_file='torchcv/datasets/kitti/kitti12_val.txt', \
transform=transform_test)
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=8,
shuffle=True,
num_workers=8)
testloader = torch.utils.data.DataLoader(testset,
batch_size=8,
shuffle=False,
num_workers=8)
# Model
print('==> Building model..')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
net = FPNSSD512(num_classes=9).to(device)
net.load_state_dict(torch.load(args.model))
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
best_loss = float('inf') # best test loss
start_epoch = 0 # start from epoch 0 or last epoch
if args.resume:
print('==> Resuming from checkpoint..')
checkpoint = torch.load(args.checkpoint)
net.load_state_dict(checkpoint['net'])
best_loss = checkpoint['loss']
start_epoch = checkpoint['epoch']
criterion = SSDLoss(num_classes=9)
from torchcv.models.fpnssd import FPNSSD512, FPNSSDBoxCoder
import numpy as np
import glob
import cv2
labelmap = [
'human', 'hold', 'stand', 'sit', 'ride', 'walk', 'look', 'hit_1', 'eat_1',
'jump', 'lay', 'talk_on_phone', 'carry', 'throw', 'catch', 'cut_1', 'run',
'work_on_computer', 'ski', 'surf', 'skateboard', 'smile', 'drink', 'kick',
'point', 'read', 'snowboard', 'eat_2', 'cut_2', 'hit_2'
]
fnt = ImageFont.truetype('Pillow/Tests/fonts/FreeMono.ttf', 40)
#print('Loading model..')
net = FPNSSD512(num_classes=31).to('cuda')
#net = torch.nn.DataParallel(net)
net.load_state_dict(torch.load('./weights/Epoch5.pth'))
#net.eval()
for image_path in glob.glob("testImages/*"):
print('Loading image..')
img = Image.open(image_path)
if img.mode != 'RGB':
img = img.convert('RGB')
ow = oh = 512
img = img.resize((ow, oh))
print('Predicting..')
transform = transforms.Compose([
transforms.ToTensor(),