import matplotlib.pyplot as plt
from torchvision import transforms
from dataset.augmentation import MaxSizeResizer, ToTensor, SquarePad, Augmenter
parser = argparse.ArgumentParser()
parser.add_argument('--input', default='data/letters/my/-11_issues.png')
parser.add_argument('--output', default='output.jpg')
parser.add_argument('--weights', default='model/letters-nandos-efficientdet-d0-e3000.pth')
parser.add_argument('--network', default='efficientdet-d0')
parser.add_argument('--num_classes', default=1)
args = parser.parse_args()
if __name__ == '__main__':
model = EfficientDet(args.num_classes, network=args.network)
model.load_state_dict(torch.load(args.weights))
model.eval()
image = cv2.imread(args.input)
transforms = transforms.Compose(
[
MaxSizeResizer(1280),
SquarePad(),
ToTensor(),
]
)
image = transforms({'image': image})['image'].float()
image = image.view(1, *image.size())
classes_, _, train_rects, output_rects = model(image.float())
out_classes, out_rects = postprocess(classes_[0], output_rects[0], 0.15)
image = image[0].int().cpu().numpy().transpose([1, 2, 0]).copy().astype('uint8')
transform=transforms.Compose([
MaxSizeResizer(1536),
SquarePad(),
ToTensor(),
]))
model = EfficientDet(train_dataset.num_classes(),
network=args.network,
device=device).to(device)
model.eval()
init_epoch = 0
if args.checkpoint:
checkpoint_name = os.path.basename(args.checkpoint)
init_epoch = int(re.findall(r'e\d+', checkpoint_name)[0][1:])
model.load_state_dict(torch.load(args.checkpoint))
model.train()
optimizer = ranger(model.parameters(), 1e-3)
for e in range(init_epoch, args.epochs):
losses = []
for images, rects, classes in DataLoader(train_dataset, 1, False):
optimizer.zero_grad()
classes_, activations, train_rects, output_rects = model(
images.float())
loss = calc_loss(train_rects, classes_, rects, classes)
losses.append(loss.item())