def infer_and_show_img(img_filepath, hmt_result, show_window=True):
"""
inference with image contains only one person
:param hmt_result:
:param img_filepath:
:return:
"""
from util.img_utils import det_landmarks
dlib_result = det_landmarks(img_filepath)
print(dlib_result)
image = cv2.imread(img_filepath)
if image.shape[0] >= image.shape[1]:
image = image[0:image.shape[1], :, :]
else:
image = image[:, 0:image.shape[0], :]
# image = cv2.resize(image, (400, 400))
h, w, c = image.shape
text_area_height = 50
final_image = np.zeros([h + text_area_height, 2 * w, c], dtype=np.uint8)
final_image[0:h, 0:w, :] = image
final_image[h:h + text_area_height, :, :] = 255
face = dlib_result[0]
if hmt_result['results']['gender'] == 'male':
color = (255, 0, 0)
else:
color = (0, 0, 255)
race_text = 'Asian' if hmt_result['results'][
'race'] == 'yellow' else 'Westerner'
cv2.rectangle(image, (face['bbox'][0], face['bbox'][1]),
(face['bbox'][2], face['bbox'][3]), color, 2)
cv2.putText(
final_image, 'Face Beauty:{0} Race:{1} Gender:{2}'.format(
str(round(hmt_result['results']['attractiveness'], 3)), race_text,
hmt_result['results']['gender']),
(int(w / 2), h + int(text_area_height / 2)), cv2.FONT_HERSHEY_SIMPLEX,
0.8, (0, 0, 255), 0, cv2.LINE_AA)
# text_color = (99, 99, 238),
for ldmk in face['landmarks']:
cv2.circle(image, (ldmk[0], ldmk[1]), 2, (255, 245, 0), -1)
final_image[0:h, w:2 * w, :] = image
save_to_dir = 'D:/TikTok_with_Annotation'
mkdirs_if_not_exist(save_to_dir)
cv2.imwrite(os.path.join(save_to_dir,
img_filepath.split(os.sep)[-1]), final_image)
if show_window:
cv2.imshow('final_image', final_image)
cv2.waitKey()
cv2.destroyAllWindows()
def video_to_imgs(video_file):
"""
write video frames to local images
:param video_file:
:return:
"""
dir_name = './{0}'.format(video_file.split('/')[-1].split('.')[0])
mkdirs_if_not_exist(dir_name)
vidcap = cv2.VideoCapture(video_file)
success, image = vidcap.read()
count = 0
while success:
cv2.imwrite(os.path.join(dir_name, "frame%d.jpg" % count), image) # save frame as JPEG file
success, image = vidcap.read()
print('Read a new frame: ', success)
count += 1
print('Images have been written successfully~~~')
def crop_faces(img_dir, type='MTCNN'):
"""
crop face region and show image window
:param img_dir:
:param type" MTCNN or dlib
:return:
"""
from mtcnn.mtcnn import MTCNN
detector = MTCNN()
fail_list = []
crop_dir = 'E:/DataSet/Face/SCUT-FBP5500/Crop'
mkdirs_if_not_exist(crop_dir)
for img_file in os.listdir(img_dir):
print('process image %s ...' % str(img_file))
if type == 'dlib':
res = det_landmarks(os.path.join(img_dir, img_file))
for i in range(len(res)):
bbox = res[i]['bbox']
image = cv2.imread(os.path.join(img_dir, img_file))
# cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (255, 0, 0), 2)
face_region = image[bbox[0]: bbox[2], bbox[1]:bbox[3]]
cv2.imwrite(os.path.join(crop_dir, img_file), face_region)
elif type == 'MTCNN':
img = cv2.imread(os.path.join(img_dir, img_file))
try:
result = detector.detect_faces(img)
# cv2.rectangle(img, (result[0]['box'][0], result[0]['box'][1]),
# (result[0]['box'][0] + result[0]['box'][2], result[0]['box'][1] + result[0]['box'][3]),
# (0, 155, 255), 2)
face_region = img[result[0]['box'][0]: result[0]['box'][0] + result[0]['box'][2],
result[0]['box'][1]: result[0]['box'][1] + result[0]['box'][3]]
cv2.imwrite(os.path.join(crop_dir, img_file), face_region)
except:
fail_list.append(img_file)
# cv2.imshow('face_region', face_region)
# cv2.waitKey()
# cv2.destroyAllWindows()
print(fail_list)
def train_model(model, train_dataloader, test_dataloader, criterion, optimizer, scheduler, num_epochs,
inference=False):
"""
train model
:param model:
:param train_dataloader:
:param test_dataloader:
:param criterion:
:param optimizer:
:param scheduler:
:param num_epochs:
:param inference:
:return:
"""
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
model = model.to(device)
if not inference:
print('Start training ResNet...')
model.train()
for epoch in range(num_epochs):
scheduler.step()
running_loss = 0.0
for i, data in enumerate(train_dataloader, 0):
images, label = data['image'], data['label']
images = images.to(device)
label = label.to(device)
optimizer.zero_grad()
pred = model(images)
loss = criterion(pred, label)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 50 == 49: # print every 50 mini-batches
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 50))
running_loss = 0.0
print('Finished training ResNet...\n')
print('Saving trained model...')
model_path_dir = './model'
mkdirs_if_not_exist(model_path_dir)
torch.save(model.state_dict(), os.path.join(model_path_dir, 'resnet.pth'))
print('ResNet has been saved successfully~')
else:
print('Loading pre-trained model...')
model.load_state_dict(torch.load(os.path.join('./model/resnet.pth')))
model.eval()
correct = 0
total = 0
with torch.no_grad():
for data in test_dataloader:
images, label = data['image'], data['label']
images = images.to(device)
label = label.to(device)
pred = model.forward(images)
_, predicted = torch.max(pred.data, 1)
total += pred.size(0)
correct += (predicted == label).sum().item()
print('Accuracy of ResNet: %f' % (correct / total))
def train_model_ft(model, dataloaders, criterion, optimizer, scheduler, num_epochs, inference):
"""
train model with fine-tune on ImageNet
:param dataloaders:
:param model:
:param criterion:
:param optimizer:
:param scheduler:
:param num_epochs:
:param inference:
:return:
"""
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
model = model.to(device)
if not inference:
since = time.time()
best_model_wts = copy.deepcopy(model.state_dict())
best_acc = 0.0
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 100)
# Each epoch has a training and validation phase
for phase in ['train', 'val']:
if phase == 'train':
scheduler.step()
model.train() # Set model to training mode
else:
model.eval() # Set model to evaluate mode
running_loss = 0.0
running_corrects = 0
# Iterate over data.
for i, data in enumerate(dataloaders[phase], 0):
inputs = data['image'].to(device)
labels = data['label'].to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward
# track history if only in train
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
# statistics
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / (dataset_sizes[phase] * batch_size)
epoch_acc = running_corrects.double() / (dataset_sizes[phase] * batch_size)
print('{} Loss: {:.4f} Acc: {:.4f}'.format(
phase, epoch_loss, epoch_acc))
# deep copy the model
if phase == 'val' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy.deepcopy(model.state_dict())
model_path_dir = './model'
mkdirs_if_not_exist(model_path_dir)
torch.save(model.state_dict(), os.path.join(model_path_dir, '{0}_Epoch_{1}.pth'.format(
model.__class__.__name__, epoch)))
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Acc: {:4f}'.format(best_acc))
# load best model weights
model.load_state_dict(best_model_wts)
else:
print('Loading pre-trained model...')
model.load_state_dict(torch.load(os.path.join('./model/%s.pth' % model.__class__.__name__)))
model.eval()
correct = 0
total = 0
y_pred = []
y_true = []
filename_list = []
probs = []
with torch.no_grad():
for data in dataloaders['test']:
images, labels, filenames = data['image'], data['label'], data['filename']
images = images.to(device)
labels = labels.to(device)
outputs = model.forward(images)
outputs = F.softmax(outputs)
# get TOP-K output labels and corresponding probabilities
topK_prob, topK_label = torch.topk(outputs, 2)
probs += topK_prob.to("cpu").detach().numpy().tolist()
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
y_pred += predicted.to("cpu").detach().numpy().tolist()
y_true += labels.to("cpu").detach().numpy().tolist()
filename_list += filenames
print('Accuracy of ResNet: %f' % (correct / total))
cm = confusion_matrix(y_true, y_pred)
print(cm)
cm = np.array(cm)
precisions = []
recalls = []
for i in range(len(cm)):
precisions.append(cm[i][i] / sum(cm[:, i].tolist()))
recalls.append(cm[i][i] / sum(cm[i, :].tolist()))
print('Precision List: ')
print(precisions)
print('Recall List: ')
print(recalls)
print("Precision of {0} on val set = {1}".format(model.__class__.__name__,
sum(precisions) / len(precisions)))
print(
"Recall of {0} on val set = {1}".format(model.__class__.__name__, sum(recalls) / len(recalls)))
print('Output CSV...')
col = ['filename', 'gt', 'pred', 'prob']
df = pd.DataFrame([[filenames[i], y_true[i], y_pred[i], probs[i][0]] for i in range(len(filenames))],
columns=col)
df.to_csv("./output-%s.csv" % model.__class__.__name__, index=False)
print('CSV has been generated...')
def train_model(model,
train_dataloader,
test_dataloader,
criterion,
optimizer,
scheduler,
num_epochs,
inference=False):
"""
train model
:param model:
:param train_dataloader:
:param test_dataloader:
:param criterion:
:param optimizer:
:param scheduler:
:param num_epochs:
:param inference:
:return:
"""
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
model = model.to(device)
if not inference:
print('Start training Model...')
model.train()
for epoch in range(num_epochs):
scheduler.step()
running_loss = 0.0
for i, data in enumerate(train_dataloader, 0):
images, label = data['image'], data['attractiveness']
images = images.to(device)
label = label.to(device).float()
optimizer.zero_grad()
pred = model(images)
loss = criterion(pred, label.unsqueeze(-1))
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 50 == 49: # print every 50 mini-batches
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 50))
running_loss = 0.0
print('Finished training Model...\n')
print('Saving trained model...')
model_path_dir = './model'
mkdirs_if_not_exist(model_path_dir)
torch.save(model.state_dict(),
os.path.join(model_path_dir, 'scutnet.pth'))
print('Model has been saved successfully~')
else:
print('Loading pre-trained model...')
model.load_state_dict(torch.load(os.path.join('./model/scutnet.pth')))
model.eval()
predicted_attractiveness_values = []
gt_attractiveness_values = []
for data in test_dataloader:
images, a_gt = data['image'], data['attractiveness']
model = model.to(device)
a_gt = a_gt.to(device)
a_pred = model.forward(images.to(device))
predicted_attractiveness_values += a_pred.to(
"cpu").data.numpy().tolist()
gt_attractiveness_values += a_gt.to("cpu").numpy().tolist()
from sklearn.metrics import mean_absolute_error, mean_squared_error
mae_lr = round(
mean_absolute_error(np.array(gt_attractiveness_values),
np.array(predicted_attractiveness_values).ravel()),
4)
rmse_lr = round(
np.math.sqrt(
mean_squared_error(
np.array(gt_attractiveness_values),
np.array(predicted_attractiveness_values).ravel())), 4)
pc = round(
np.corrcoef(np.array(gt_attractiveness_values),
np.array(predicted_attractiveness_values).ravel())[0, 1],
4)
print(
'===============The Mean Absolute Error of Model is {0}===================='
.format(mae_lr))
print(
'===============The Root Mean Square Error of Model is {0}===================='
.format(rmse_lr))
print(
'===============The Pearson Correlation of Model is {0}===================='
.format(pc))
def train_model_ft(model, dataloaders, criterion, optimizer, scheduler,
num_epochs, inference):
"""
train model with fine-tune on ImageNet
:param dataloaders:
:param model:
:param criterion:
:param optimizer:
:param scheduler:
:param num_epochs:
:param inference:
:return:
"""
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
model = model.to(device)
if not inference:
since = time.time()
best_model_wts = copy.deepcopy(model.state_dict())
best_acc = 0.0
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 100)
# Each epoch has a training and validation phase
for phase in ['train', 'val']:
if phase == 'train':
scheduler.step()
model.train() # Set model to training mode
else:
model.eval() # Set model to evaluate mode
running_loss = 0.0
running_corrects = 0
# Iterate over data.
for inputs, labels in dataloaders[phase]:
inputs = inputs.to(device)
labels = labels.to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward
# track history if only in train
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
# statistics
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / (dataset_sizes[phase] *
cfg['batch_size'])
epoch_acc = running_corrects.double() / (dataset_sizes[phase] *
cfg['batch_size'])
print('{} Loss: {:.4f} Acc: {:.4f}'.format(
phase, epoch_loss, epoch_acc))
# deep copy the model
if phase == 'val' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy.deepcopy(model.state_dict())
model_path_dir = './model'
mkdirs_if_not_exist(model_path_dir)
torch.save(
model.state_dict(),
os.path.join(
model_path_dir, '{0}_Epoch_{1}.pth'.format(
model.__class__.__name__, epoch)))
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Acc: {:4f}'.format(best_acc))
# load best model weights
model.load_state_dict(best_model_wts)
else:
print('Loading pre-trained model...')
model.load_state_dict(
torch.load(
os.path.join('./model/%s.pth' % model.__class__.__name__)))
model.eval()
correct = 0
total = 0
with torch.no_grad():
for data in dataloaders['test']:
images, label = data['image'], data['label']
images = images.to(device)
label = label.to(device)
pred = model.forward(images)
_, predicted = torch.max(pred.data, 1)
total += pred.size(0)
correct += (predicted == label).sum().item()
print('Accuracy of ResNet: %f' % (correct / total))
def train_hmtnet(hmt_net,
train_loader,
test_loader,
num_epochs,
inference=False):
"""
train HMT-Net
:param hmt_net:
:param train_loader:
:param test_loader:
:param num_epochs:
:param inference:
:return:
"""
print(hmt_net)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
print("We are running on", torch.cuda.device_count(), "GPUs!")
hmt_net = nn.DataParallel(hmt_net)
hmt_net = hmt_net.to(device)
criterion = HMTLoss()
optimizer = optim.SGD(hmt_net.parameters(),
lr=0.001,
momentum=0.9,
weight_decay=5e-2)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)
if not inference:
hmt_net.train()
for epoch in range(num_epochs): # loop over the dataset multiple times
exp_lr_scheduler.step()
running_loss = 0.0
for i, data in enumerate(train_loader, 0):
# get the inputs
inputs, gender, race, attractiveness = data['image'], data['gender'], data['race'], \
data['attractiveness']
inputs, gender, race, attractiveness = inputs.to(
device), gender.to(device), race.to(
device), attractiveness.float().to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
g_pred, r_pred, a_pred = hmt_net.forward(inputs)
g_pred = g_pred.view(cfg['batch_size'], -1)
r_pred = r_pred.view(cfg['batch_size'], -1)
a_pred = a_pred.view(cfg['batch_size'])
loss = criterion(g_pred, gender, r_pred, race, a_pred,
attractiveness)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 50 == 49: # print every 100 mini-batches
print('[%d, %5d] loss: %.5f' %
(epoch + 1, i + 1, running_loss / 50))
running_loss = 0.0
hmt_net.eval()
predicted_attractiveness_values = []
gt_attractiveness_values = []
total = 0
g_correct = 0
r_correct = 0
for data in test_loader:
images, g_gt, r_gt, a_gt = data['image'], data['gender'], data['race'], \
data['attractiveness']
hmt_net = hmt_net.to(device)
g_gt = g_gt.to(device)
r_gt = r_gt.to(device)
a_gt = a_gt.to(device)
images = images.to(device)
bs, ncrops, c, h, w = images.size()
g_pred, r_pred, a_pred = hmt_net(images.view(
-1, c, h, w)) # fuse batch size and ncrops
a_pred = a_pred.view(bs, ncrops, -1).mean(1) # avg over crops
g_pred = g_pred.view(bs, ncrops, -1).mean(1) # avg over crops
r_pred = r_pred.view(bs, ncrops, -1).mean(1) # avg over crops
predicted_attractiveness_values += a_pred.to(
"cpu").data.numpy().tolist()
gt_attractiveness_values += a_gt.to("cpu").numpy().tolist()
g_pred = g_pred.view(cfg['batch_size'], -1)
r_pred = r_pred.view(cfg['batch_size'], -1)
_, g_predicted = torch.max(g_pred.data, 1)
_, r_predicted = torch.max(r_pred.data, 1)
total += g_gt.size(0)
g_correct += (g_predicted == g_gt).sum().item()
r_correct += (r_predicted == r_gt).sum().item()
print('total = %d ...' % total)
print('Gender correct sample = %d ...' % g_correct)
print('Race correct sample = %d ...' % r_correct)
print('Accuracy of Race Classification: %.4f' %
(r_correct / total))
print('Accuracy of Gender Classification: %.4f' %
(g_correct / total))
mae_lr = round(
mean_absolute_error(
np.array(gt_attractiveness_values),
np.array(predicted_attractiveness_values).ravel()), 4)
rmse_lr = round(
np.math.sqrt(
mean_squared_error(
np.array(gt_attractiveness_values),
np.array(predicted_attractiveness_values).ravel())), 4)
pc = round(
np.corrcoef(
np.array(gt_attractiveness_values),
np.array(predicted_attractiveness_values).ravel())[0, 1],
4)
print(
'===============The Mean Absolute Error of HMT-Net is {0}===================='
.format(mae_lr))
print(
'===============The Root Mean Square Error of HMT-Net is {0}===================='
.format(rmse_lr))
print(
'===============The Pearson Correlation of HMT-Net is {0}===================='
.format(pc))
model_path_dir = './model'
file_utils.mkdirs_if_not_exist(model_path_dir)
torch.save(
hmt_net.state_dict(),
os.path.join(model_path_dir,
'hmt-net-fbp-epoch-%d.pth' % (epoch + 1)))
hmt_net.train()
print('Finished Training')
print('Save trained model...')
model_path_dir = './model'
file_utils.mkdirs_if_not_exist(model_path_dir)
torch.save(hmt_net.state_dict(),
os.path.join(model_path_dir, 'hmt-net-fbp.pth'))
else:
print('Loading pre-trained model...')
hmt_net.load_state_dict(
torch.load(os.path.join('./model/hmt-net-fbp.pth')))
hmt_net.eval()
predicted_attractiveness_values = []
gt_attractiveness_values = []
total = 0
g_correct = 0
r_correct = 0
for data in test_loader:
images, g_gt, r_gt, a_gt = data['image'], data['gender'], data['race'], \
data['attractiveness']
hmt_net = hmt_net.to(device)
g_gt = g_gt.to(device)
r_gt = r_gt.to(device)
a_gt = a_gt.to(device)
images = images.to(device)
bs, ncrops, c, h, w = images.size()
g_pred, r_pred, a_pred = hmt_net(images.view(
-1, c, h, w)) # fuse batch size and ncrops
a_pred = a_pred.view(bs, ncrops, -1).mean(1) # avg over crops
g_pred = g_pred.view(bs, ncrops, -1).mean(1) # avg over crops
r_pred = r_pred.view(bs, ncrops, -1).mean(1) # avg over crops
predicted_attractiveness_values += a_pred.to(
"cpu").data.numpy().tolist()
gt_attractiveness_values += a_gt.to("cpu").numpy().tolist()
# g_pred = g_pred.view(-1, g_pred.numel())
# r_pred = r_pred.view(-1, r_pred.numel())
g_pred = g_pred.view(cfg['batch_size'], -1)
r_pred = r_pred.view(cfg['batch_size'], -1)
_, g_predicted = torch.max(g_pred.data, 1)
_, r_predicted = torch.max(r_pred.data, 1)
total += g_gt.size(0)
g_correct += (g_predicted == g_gt).sum().item()
r_correct += (r_predicted == r_gt).sum().item()
print('total = %d ...' % total)
print('Gender correct sample = %d ...' % g_correct)
print('Race correct sample = %d ...' % r_correct)
print('Accuracy of Race Classification: %.4f' % (r_correct / total))
print('Accuracy of Gender Classification: %.4f' % (g_correct / total))
mae_lr = round(
mean_absolute_error(np.array(gt_attractiveness_values),
np.array(predicted_attractiveness_values).ravel()),
4)
rmse_lr = round(
np.math.sqrt(
mean_squared_error(
np.array(gt_attractiveness_values),
np.array(predicted_attractiveness_values).ravel())), 4)
pc = round(
np.corrcoef(np.array(gt_attractiveness_values),
np.array(predicted_attractiveness_values).ravel())[0, 1],
4)
print(
'===============The Mean Absolute Error of HMT-Net is {0}===================='
.format(mae_lr))
print(
'===============The Root Mean Square Error of HMT-Net is {0}===================='
.format(rmse_lr))
print(
'===============The Pearson Correlation of HMT-Net is {0}===================='
.format(pc))
def train_anet(model_ft,
train_loader,
test_loader,
criterion,
num_epochs=200,
inference=False):
"""
train ANet
:param model_ft:
:param train_loader:
:param test_loader:
:param criterion:
:param num_epochs:
:param inference:
:return:
"""
num_ftrs = model_ft.fc8.in_channels
model_ft.fc8 = nn.Conv2d(num_ftrs, 1, 1)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
print("We are running on", torch.cuda.device_count(), "GPUs!")
model_ft = nn.DataParallel(model_ft)
model_ft = model_ft.to(device)
optimizer_ft = optim.SGD(model_ft.parameters(),
lr=0.001,
momentum=0.9,
weight_decay=1e-4)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft,
step_size=50,
gamma=0.1)
if not inference:
for epoch in range(num_epochs): # loop over the dataset multiple times
exp_lr_scheduler.step()
model_ft.train()
running_loss = 0.0
for i, data in enumerate(train_loader, 0):
# get the inputs
inputs, labels = data['image'], data['attractiveness']
model_ft = model_ft.to(device)
inputs, labels = inputs.to(device), labels.float().to(device)
# zero the parameter gradients
optimizer_ft.zero_grad()
# forward + backward + optimize
outputs = model_ft(inputs)
outputs = outputs.view(cfg['batch_size'])
loss = criterion(outputs, labels)
loss.backward()
optimizer_ft.step()
# print statistics
running_loss += loss.item()
if i % 50 == 49: # print every 50 mini-batches
print('[%d, %5d] loss: %.5f' %
(epoch + 1, i + 1, running_loss / 50))
running_loss = 0.0
print('Finished Training')
print('Save trained model...')
model_path_dir = './model'
file_utils.mkdirs_if_not_exist(model_path_dir)
torch.save(model_ft.state_dict(),
os.path.join(model_path_dir, 'anet.pth'))
else:
print('Loading pre-trained model...')
model_ft.load_state_dict(torch.load(os.path.join('./model/anet.pth')))
model_ft.eval()
predicted_labels = []
gt_labels = []
for i, data in enumerate(test_loader, 0):
images, labels = data['image'], data['attractiveness']
model_ft = model_ft.to(device)
labels = labels.to(device)
outputs = model_ft.forward(images.to(device))
predicted_labels += outputs.to("cpu").data.numpy().tolist()
gt_labels += labels.to("cpu").numpy().tolist()
from sklearn.metrics import mean_absolute_error, mean_squared_error
mae_lr = round(
mean_absolute_error(np.array(gt_labels),
np.array(predicted_labels).ravel()), 4)
rmse_lr = round(
np.math.sqrt(
mean_squared_error(np.array(gt_labels),
np.array(predicted_labels).ravel())), 4)
pc = round(
np.corrcoef(np.array(gt_labels),
np.array(predicted_labels).ravel())[0, 1], 4)
print(
'===============The Mean Absolute Error of ANet is {0}===================='
.format(mae_lr))
print(
'===============The Root Mean Square Error of ANet is {0}===================='
.format(rmse_lr))
print(
'===============The Pearson Correlation of ANet is {0}===================='
.format(pc))
def train_gnet(model,
train_loader,
test_loader,
criterion,
optimizer,
num_epochs=200,
inference=False):
"""
train GNet
:param model:
:param train_loader:
:param test_loader:
:param criterion:
:param optimizer:
:param num_epochs:
:return:
"""
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
print("We are running on", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
if not inference:
exp_lr_scheduler.step()
model.train()
for epoch in range(num_epochs): # loop over the dataset multiple times
running_loss = 0.0
for i, data in enumerate(train_loader, 0):
# get the inputs
# inputs, labels = data
inputs, labels = data['image'], data['gender']
model = model.to(device)
inputs, labels = inputs.to(device), labels.to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = model.forward(inputs)
outputs = outputs.view(cfg['batch_size'], -1)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 50 == 49: # print every 50 mini-batches
print('[%d, %5d] loss: %.5f' %
(epoch + 1, i + 1, running_loss / 50))
running_loss = 0.0
print('Finished Training')
print('Save trained model...')
model_path_dir = './model'
file_utils.mkdirs_if_not_exist(model_path_dir)
torch.save(model.state_dict(), os.path.join(model_path_dir,
'gnet.pth'))
else:
print('Loading pre-trained model...')
model.load_state_dict(torch.load(os.path.join('./model/gnet.pth')))
model.eval()
correct = 0
total = 0
for data in test_loader:
# images, labels = data
images, labels = data['image'], data['gender']
model = model.to(device)
labels = labels.to(device)
outputs = model.forward(images.to(device))
outputs = outputs.view(cfg['batch_size'], 2)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('correct = %d ...' % correct)
print('total = %d ...' % total)
print('Accuracy of the network on test images: %f' % (correct / total))
def finetune_vgg_m_model(model_ft,
train_loader,
test_loader,
criterion,
num_epochs=200,
inference=False):
"""
fine-tune VGG M Face Model
:param model_ft:
:param train_loader:
:param test_loader:
:param criterion:
:param num_epochs:
:param inference:
:return:
"""
num_ftrs = model_ft.fc8.in_channels
model_ft.fc8 = nn.Conv2d(num_ftrs, 2, 1)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
print("We are running on", torch.cuda.device_count(), "GPUs!")
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model_ft = nn.DataParallel(model_ft)
model_ft = model_ft.to(device)
optimizer_ft = optim.SGD(model_ft.parameters(),
lr=0.001,
momentum=0.9,
weight_decay=1e-4)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft,
step_size=50,
gamma=0.1)
if not inference:
for epoch in range(num_epochs): # loop over the dataset multiple times
exp_lr_scheduler.step()
model_ft.train()
running_loss = 0.0
for i, data in enumerate(train_loader, 0):
# get the inputs
# inputs, labels = data
inputs, labels = data['image'], data['attractiveness'].float()
model_ft = model_ft.to(device)
inputs, labels = inputs.to(device), labels.to(device)
# zero the parameter gradients
optimizer_ft.zero_grad()
# forward + backward + optimize
outputs = model_ft.forward(inputs)
outputs = (torch.sum(outputs, dim=1) / 2).view(
cfg['batch_size'], 1)
outputs = outputs.view(cfg['batch_size'])
loss = criterion(outputs, labels)
loss.backward()
optimizer_ft.step()
# print statistics
running_loss += loss.item()
if i % 50 == 49: # print every 50 mini-batches
print('[%d, %5d] loss: %.5f' %
(epoch + 1, i + 1, running_loss / 50))
running_loss = 0.0
print('Finished Training')
print('Save trained model...')
model_path_dir = './model'
file_utils.mkdirs_if_not_exist(model_path_dir)
torch.save(model_ft.state_dict(),
os.path.join(model_path_dir, 'ft_vgg_m.pth'))
else:
print('Loading pre-trained model...')
model_ft.load_state_dict(
torch.load(os.path.join('./model/ft_vgg_m.pth')))
model_ft.eval()
correct = 0
total = 0
# for data in test_loader:
for i, data in enumerate(test_loader, 0):
# images, labels = data
images, labels = data['image'], data['attractiveness'].float()
model_ft = model_ft.to(device)
labels = labels.to(device)
outputs = model_ft.forward(images.to(device))
outputs = outputs.view(cfg['batch_size'])
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('correct = %d ...' % correct)
print('total = %d ...' % total)
print('Accuracy of the network on the test images: %f' % (correct / total))