def classify_single_image(image_path): # model loading model = resnet.ResNet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], NUM_CLASSES) model.eval() if USE_GPU: cudnn.enabled = True softmax = nn.Softmax().cuda() model.cuda() saved_state_dict = torch.load(MODEL_PATH) else: softmax = nn.Softmax() saved_state_dict = torch.load(MODEL_PATH, map_location='cpu') load_filtered_state_dict(model, saved_state_dict, ignore_layer=[], reverse=False, gpu=cudnn.enabled) transformations = transforms.Compose( [transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)), transforms.ToTensor()]) imgs = torch.FloatTensor(1, 3, IMAGE_SIZE, IMAGE_SIZE) if USE_GPU: imgs = imgs.cuda() #image loading imgs[0] = transformations(Image.open(image_path).convert("RGB")) pred = model(imgs) pred = softmax(pred) print(pred.cpu().detach().numpy()) _, pred_1 = pred.topk(1, 1, True, True) c = default_class[pred_1.cpu().numpy()[0][0]] print("{} -- {}".format(image_path, c))
def main(args): model = resnet.ResNet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], args.num_classes) saved_state_dict = torch.load(args.saved_model) transformations = transforms.Compose([ transforms.Resize((args.image_size, args.image_size)), transforms.ToTensor() ]) if args.gpu[0] >= 0: cudnn.enabled = True softmax = nn.Softmax().cuda() model.cuda() else: softmax = nn.Softmax() load_filtered_state_dict(model, saved_state_dict, ignore_layer=[], reverse=True) test_x, test_y, classes_names = get_dataset(args.test_data_dir) test_dataset = DataWrapper(test_x, test_y, transformations, augumentation=False) test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=1) classes, indices = np.unique(test_y, return_index=True) #n = (test_dataset.__len__() + args.batch_size - 1) / args.batch_size * args.batch_size n = test_dataset.__len__() y_pred = np.zeros((n)) y = np.zeros((n)) count = 0 for i, (images, labels, names) in enumerate(test_loader): images = Variable(images) labels = Variable(labels) if args.gpu[0] >= 0: images = images.cuda() labels = labels.cuda() label_pred = model(images) label_pred = softmax(label_pred) n = images.size()[0] _, label_pred = label_pred.topk(1, 1, True, True) y_pred[count:count + n] = label_pred.view(-1).cpu().numpy() y[count:count + n] = labels.data.cpu().numpy() count += n plot(y, y_pred, classes_names)
def main(args): cudnn.enabled = True print('Loading data.') transformations = transforms.Compose([ transforms.Resize(240), transforms.RandomCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) train_x, train_y, classes_names = get_dataset(args.trainning_data_dir) test_x, test_y, _ = get_dataset(args.validation_data_dir) num_classes = len(classes_names) trainning_dataset = DataWrapper(train_x, train_y, transformations) eval_dataset = DataWrapper(test_x, test_y, transformations) train_loader = torch.utils.data.DataLoader(dataset=trainning_dataset, batch_size=args.batch_size, shuffle=True, num_workers=16) eval_loader = torch.utils.data.DataLoader(dataset=eval_dataset, batch_size=args.batch_size, shuffle=True, num_workers=16) n = trainning_dataset.__len__() print(n) # ResNet50 structure model = resnet.ResNet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], num_classes) if args.saved_model: print('Loading model.') saved_state_dict = torch.load(args.saved_model) # 'origin model from pytorch' if 'resnet' in args.saved_model: load_filtered_state_dict(model, saved_state_dict, ignore_layer=[], reverse=False) else: load_filtered_state_dict(model, saved_state_dict, ignore_layer=[]) crossEntropyLoss = nn.CrossEntropyLoss().cuda() softmax = nn.Softmax().cuda() optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) # multi-gpu model = nn.DataParallel(model, device_ids=[0, 1]) model.cuda() Save_model = SaveBestModel(save_dir=args.save_path) Writer = SummaryWriter() step = 0 for epoch in range(args.num_epochs): evaluate(eval_loader, model, Writer, step, Save_model, epoch) step = train(train_loader, model, crossEntropyLoss, optimizer, Writer, args.batch_size, epoch, step, n)
def main(args): model = resnet.ResNet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], args.num_classes) #model = dpn92(num_classes=args.num_classes) transformations = transforms.Compose([ transforms.Resize((args.image_size, args.image_size)), transforms.ToTensor() ]) if args.gpu[0] >= 0: cudnn.enabled = True softmax = nn.Softmax().cuda() model.cuda() saved_state_dict = torch.load(args.saved_model) else: softmax = nn.Softmax() saved_state_dict = torch.load(args.saved_model, map_location='cpu') load_filtered_state_dict(model, saved_state_dict, ignore_layer=[], reverse=True, gpu=cudnn.enabled) imgs_path = glob.glob(os.path.join(args.test_data_dir, '*.jpg')) imgs_path += glob.glob(os.path.join(args.test_data_dir, '*.jpeg')) imgs_path += glob.glob(os.path.join(args.test_data_dir, '*.png')) for i in xrange((len(imgs_path) + args.batch_size - 1) / args.batch_size): if args.gpu[0] >= 0: imgs = torch.FloatTensor(args.batch_size, 3, args.image_size, args.image_size).cuda() else: imgs = torch.FloatTensor(args.batch_size, 3, args.image_size, args.image_size) for j in xrange(min(args.batch_size, len(imgs_path))): img = Image.open(imgs_path[i * args.batch_size + j]) img = img.convert("RGB") imgs[j] = transformations(img) pred = model(imgs) pred = softmax(pred) _, pred_1 = pred.topk(1, 1, True, True) for j in xrange(min(args.batch_size, len(imgs_path))): c = default_class[pred_1.cpu().numpy()[0][0]] print("{} -- {} {}".format(imgs_path[i * args.batch_size + j], pred_1, c)) if args.save_path: img_numpy = imgs[j].cpu().numpy() img_numpy = img_numpy * 255 # change to channel last img_numpy = np.transpose(img_numpy, (1, 2, 0)).astype(np.uint8) # rgb to bgr img_numpy = img_numpy[..., ::-1].copy() cv2.putText(img_numpy, c, (20, 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA) save_path = os.path.join( args.save_path, os.path.basename(imgs_path[i * args.batch_size + j])) cv2.imwrite(save_path, img_numpy)
# model loading print("model loading start") model = resnet.ResNet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], NUM_CLASSES) model.eval() if USE_GPU: cudnn.enabled = True softmax = nn.Softmax().cuda() model.cuda() saved_state_dict = torch.load(MODEL_PATH) else: softmax = nn.Softmax() saved_state_dict = torch.load(MODEL_PATH, map_location='cpu') load_filtered_state_dict(model, saved_state_dict, ignore_layer=[], reverse=False, gpu=cudnn.enabled) transformations = transforms.Compose( [transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)), transforms.ToTensor()]) print("model loading end") from flask import Flask import json import os app = Flask(__name__)