generate_db.py

import cv2
import numpy as np
import torch
import dlib
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import ToTensor

import utils.sqlrequest as sql
import utils.align_dlib as align
from utils.openface import prepareOpenFace

# Yep, this is hardcoded for now.
# TODO: Flags.
import os
import time
os.environ['CUDA_VISIBLE_DEVICES'] = str(2)

start_time = time.time()
def tprint(*args, **kwargs):
    t = time.time() - start_time
    print('[%5.3f sec]'%t, *args, **kwargs)

tprint('Initializing SQL database connection...\n', sql.initDB())
def cvLoad(imgbytes):
    img = np.frombuffer(imgbytes, dtype=np.uint8)
    img = cv2.imdecode(img, cv2.IMREAD_COLOR)
    return img

result = sql.db_query('SELECT full_frame, xmin, ymin, xmax, ymax, card_ID FROM photo WHERE is_it_sure = 1', 
                      as_dict=False)

tprint('#Entries:', len(result))
aligner = align.AlignDlib('weights/shape_predictor_68_face_landmarks.dat')
def alignImage(entry):
    imgbytes, xmin, ymin, xmax, ymax = entry[:5]
    img = cvLoad(imgbytes)[:, :, ::-1]
    bb = dlib.rectangle(xmin, ymin, xmax, ymax)
    img_aligned = aligner.align(96, img, bb)
    
    return img_aligned

class AlignedFaceDataset(Dataset):
    
    def __init__(self, sql_entries):
        self.entries = sql_entries
        self.tensor_converter = ToTensor()
        
    def __len__(self):
        return len(self.entries)
    
    def __getitem__(self, idx):
        img = alignImage(self.entries[idx])
        return self.tensor_converter(img)

'''    
tprint('Aligning faces...')
aligned_imgs = []
for i, entry in enumerate(result):
    aligned_imgs.append(alignImage(result[0]))
    if i % 1000 == 0:
        tprint('[%5d/%5d]'%(i, len(result)))
'''
tprint('Initializing network...')
dataset = AlignedFaceDataset(result)
net = prepareOpenFace(gpuDevice=0)
weights = torch.load('weights/openface.pth')
net.load_state_dict(weights)
_ = net.cuda()


tprint('Inferring aligned images...')
embedded_db = []
batch_size = 1024
dataloader = DataLoader(dataset, batch_size=batch_size,
                        shuffle=False, num_workers=35)

for i, x in enumerate(dataloader):
    input_var = torch.autograd.Variable(
        x, requires_grad=False, volatile=True).cuda()
    
    output_var, _ = net(input_var)
    embedded_db.append(output_var.data)
    tprint('[%5d/%5d]'%(i*batch_size, len(dataset)))
    
'''
batch_start_idx = 0
while batch_start_idx < len(aligned_imgs):
    input_tensor = [tensor_converter(x) for x in 
                    aligned_imgs[batch_start_idx:batch_start_idx+batch_size]]
    
    input_var = torch.autograd.Variable(
        torch.stack(input_tensor), 
        requires_grad=False, volatile=True).cuda()
    
    output_var, _ = net(input_var)
    embedded_db.append(output_var.data)
    tprint('[%5d/%5d]'%(batch_start_idx, len(aligned_imgs)))
    batch_start_idx += batch_size
'''
embedded = torch.cat(embedded_db)
embedded_data = embedded.cpu()
unzipped_result = list(zip(*result))
card_id = unzipped_result[-1]
deploy_database = {
    'id': card_id,
    'emb': embedded_data
}
torch.save(deploy_database, 'AUTOGENERATED-DB.tar')
tprint('Done')