import os
import torch
import cv2
from crnn import CRNN
from tqdm import tqdm
import csv
import numpy as np
model = CRNN()
model.load_state_dict(torch.load('55acc.pt'))
model.eval()
model.to('cuda')
data_dir = "qia2020/test/"
emo = {0: 'hap', 1: 'sur', 2: 'neu', 3: 'fea', 4: 'dis', 5: 'ang', 6: 'sad'}
with open('test_confirm.csv', 'w', newline='') as csvfile:
writer = csv.writer(csvfile)
writer.writerow(['FileID', 'Emotion'])
for filename in tqdm(sorted(os.listdir(data_dir))):
if not filename.endswith(".mp4"):
continue
f = 'torch_video_3_test/' + filename[:5] + '.pt'
X = torch.load(f)
X = X.unsqueeze(0).to('cuda:0')
with np.load(data_dir + filename[:5] + '.npz') as data:
T = torch.Tensor(data['word_embed'])
def train(root,
start_epoch,
epoch_num,
letters,
net=None,
lr=0.1,
fix_width=True):
"""
Train CRNN model
Args:
root (str): Root directory of dataset
start_epoch (int): Epoch number to start
epoch_num (int): Epoch number to train
letters (str): Letters contained in the data
net (CRNN, optional): CRNN model (default: None)
lr (float, optional): Coefficient that scale delta before it is applied
to the parameters (default: 1.0)
fix_width (bool, optional): Scale images to fixed size (default: True)
Returns:
CRNN: Trained CRNN model
"""
# load data
trainloader = load_data(root, training=True, fix_width=fix_width)
if not net:
# create a new model if net is None
net = CRNN(1, len(letters) + 1)
criterion = torch.nn.CTCLoss()
optimizer = optim.Adadelta(net.parameters(), lr=lr, weight_decay=1e-3)
# use gpu or not
use_cuda = torch.cuda.is_available()
use_cuda = False
device = torch.device('cuda' if use_cuda else 'cpu')
if use_cuda:
net = net.to(device)
criterion = criterion.to(device)
else:
print("***** Warning: Cuda isn't available! *****")
# get encoder and decoder
labeltransformer = LabelTransformer(letters)
print('==== Training.. ====')
# .train() has any effect on Dropout and BatchNorm.
net.train()
for epoch in range(start_epoch, start_epoch + epoch_num):
print('---- epoch: %d ----' % (epoch, ))
loss_sum = 0
for i, (img, label) in enumerate(trainloader):
label, label_length = labeltransformer.encode(label)
img = img.to(device)
optimizer.zero_grad()
# put images in
outputs = net(img)
output_length = torch.IntTensor([outputs.size(0)] *
outputs.size(1))
# calc loss
loss = criterion(outputs, label, output_length, label_length)
# update
loss.backward()
optimizer.step()
loss_sum += loss.item()
print('loss = %f' % (loss_sum, ))
print('Finished Training')
return net
