Example #1
0
def main():
    net = Net().to(setting.device)
    net.eval()
    net.load_state_dict(torch.load('model.pt'))
    print('model has been loaded.')
    correct = 0
    valid_dataloader = data_preprocess.get_valid_dataloader()
    total = len(os.listdir(setting.valid_folder_path))
    with torch.no_grad():
        
        miss_character = {}
        for (imgs, labels) in tqdm((valid_dataloader)):
            imgs, labels = imgs.to(setting.device), labels.to(setting.device)
            labels_ohe_predict = net(imgs)
            # for each img in one batch
            for single in range(labels_ohe_predict.shape[0]):              
                single_labels_ohe_predict = labels_ohe_predict[single, :]
                predict_label = ''
                # get predict_label
                for slice in range(setting.char_num):
                    char = ohe.num2char[np.argmax(
                        single_labels_ohe_predict[slice*setting.pool_length:(slice+1)*setting.pool_length].cpu().data.numpy())]
                    predict_label += char
                # get true label
                true_label = ohe.decode(labels[single, :].cpu().numpy())
                # print('true label:', true_label, '   predict label:', predict_label)
                if predict_label == true_label:
                    correct += 1
                else:
                    
                    for i in range(setting.char_num):
                        if predict_label[i] != true_label[i]:
                            error_info = '{} -> {}'.format(true_label[i], predict_label[i])
                            if error_info in miss_character:
                                miss_character[error_info] +=1
                            else:
                                miss_character[error_info] =1
    sorted_miss = sorted(miss_character.items(), key=lambda kv:kv[1], reverse=True)
    sorted_miss=collections.OrderedDict(sorted_miss)            
    with open('miss_character.txt','w') as f:
        for i in sorted_miss:
            f.write('{} : {}\n'.format(i, sorted_miss[i]))
    print('accuracy: {}/{} -- {:.4f}'.format(correct, total, correct/total))
animals_dataset = AnimalsDataset(filename='meta-data_new/meta-data/train.csv', root_dir='./train_new/train/', train=True,
                                 transform=data_transform1)
test_dataset = AnimalsDataset(filename='meta-data_new/meta-data/test.csv', root_dir='./test_new/test', train=False,
                              transform=data_transform1)

test_loader = DataLoader(test_dataset, batch_size=60)

use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
print(device)

net = Net().to(device)
print(net)
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

net.load_state_dict(torch.load("./models/model_ep50.net"))

net.eval()

#######################
# Testing an image
data_iter = iter(test_dataset)
sample = next(data_iter)
data, labels = sample['image'], sample['labels']

img = data
img = torch.unsqueeze(img, 0)
img = img.float().to(device)

out = net(img)
prediction = torch.nn.functional.softmax(out)

def process_img(original_image):
    processed_img = cv2.cvtColor(original_image, cv2.COLOR_BGR2GRAY)
    #processed_img = cv2.Canny(processed_img, threshold1 = 100, threshold2 = 200)
    return processed_img


def image_to_tensor(img):
    img = cv2.resize(img, dsize=(48, 48), interpolation=cv2.INTER_CUBIC)
    return torch.Tensor(img.reshape(1, 1, 48, 48))


model = Net()
dir_path = os.path.dirname(os.path.realpath(__file__))
model.load_state_dict(torch.load(dir_path + '/models/current_model.pth'))

curr_emotion = "Neutral"
emotion_ave = []
last_time = time.time()

while (True):
    #Grabs image from screen
    screen = np.array(ImageGrab.grab(bbox=(300, 300, 500, 500)))
    processed_img = cv2.cvtColor(screen, cv2.COLOR_BGR2RGB)
    new_screen = process_img(screen)

    # Convert iamge to tensor and predict the emotion
    x = new_screen
    x = image_to_tensor(x)
    output = model(x)