def main(): net = GazeNet() net = DataParallel(net) net.cuda() # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") pretrained_dict = torch.load('../model/trained_model.pkl') model_dict = net.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } # print(pretrained_dict.summary()) model_dict.update(pretrained_dict) net.load_state_dict(model_dict) test_image_path = sys.argv[1] xi, yi = detect_head(test_image_path) x = float(xi) y = float(yi) #x = float(sys.argv[2]) #y = float(sys.argv[3]) heatmap, p_x, p_y = test(net, test_image_path, (x, y)) draw_result(test_image_path, (x, y), heatmap, (p_x, p_y)) print(p_x, p_y)
def main(): test_set = GazeDataset( root_dir= '/home/emannuell/Documentos/mestrado/GazeFollowing/GazeFollowData/data_new/', mat_file='', training='test') test_data_loader = DataLoader(test_set, batch_size=1, shuffle=False, num_workers=8) net = GazeNet() net = DataParallel(net) net.cuda() pretrained_dict = torch.load( '/home/emannuell/Documentos/mestrado/GazeFollowing/model/test04/epoch_15_loss_0.0558342523873.pkl' ) model_dict = net.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } model_dict.update(pretrained_dict) net.load_state_dict(model_dict) test(net, test_data_loader)
def home(): #python3 inference.py ../images/00004844.jpg 0.35636 0.23724 net = GazeNet() net = DataParallel(net) net.cuda() pretrained_dict = torch.load('../model/trained_model.pkl') model_dict = net.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } model_dict.update(pretrained_dict) net.load_state_dict(model_dict) if request.method == 'POST': f = request.files['file'] f.save(secure_filename(f.filename)) fname = f.filename print(fname) path2 = '.' test_image_path = os.path.join(path2, f.filename) xi, yi = detect_head(test_image_path) x = float(xi) y = float(yi) print(test_image_path) # x = 0.372 # y = 0.22267 # 0.372,0.22267 heatmap, p_x, p_y = test(net, test_image_path, (x, y)) resultimg = draw_result(test_image_path, (x, y), heatmap, (p_x, p_y)) #print(resultimg) img = Image.fromarray(resultimg, 'RGB') #print(img) # path = '/content/drive/My Drive/AI_PROJECT/GazeFollowing/code/static' # cv2.imwrite(os.path.join(path2 , 'tmp.jpg'), img) print(p_x, p_y) outim = ['tmp.png'] for o in outim: shutil.copy(o, './static') return render_template('index.html')
def main(): test_set = GazeDataset(root_dir='../GazeFollowData/', mat_file='../GazeFollowData/test2_annotations.mat', training='test') test_data_loader = DataLoader(test_set, batch_size=1, shuffle=False, num_workers=8) net = GazeNet() net = DataParallel(net) net.cuda() pretrained_dict = torch.load('../model/pretrained_model.pkl') model_dict = net.state_dict() pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} model_dict.update(pretrained_dict) net.load_state_dict(model_dict) test(net, test_data_loader)
def main(): net = GazeNet() net = DataParallel(net) net.cuda() pretrained_dict = torch.load('../model/pretrained_model.pkl') model_dict = net.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } model_dict.update(pretrained_dict) net.load_state_dict(model_dict) test_image_path = sys.argv[1] x = float(sys.argv[2]) y = float(sys.argv[3]) heatmap, p_x, p_y = test(net, test_image_path, (x, y)) draw_result(test_image_path, (x, y), heatmap, (p_x, p_y)) print(p_x, p_y)
from scipy import signal from utils import data_transforms from utils import get_paste_kernel, kernel_map net = GazeNet() net = DataParallel(net) net.cpu() torch.device('cpu') pretrained_dict = torch.load('./savedmodels/pretrained_model.pkl', map_location=lambda storage, loc: storage) model_dict = net.state_dict() pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} model_dict.update(pretrained_dict) net.load_state_dict(model_dict) def generate_data_field(eye_point): """eye_point is (x, y) and between 0 and 1""" height, width = 224, 224 x_grid = np.array(range(width)).reshape([1, width]).repeat(height, axis=0) y_grid = np.array(range(height)).reshape([height, 1]).repeat(width, axis=1) grid = np.stack((x_grid, y_grid)).astype(np.float32) x, y = eye_point x, y = x * width, y * height grid -= np.array([x, y]).reshape([2, 1, 1]).astype(np.float32) norm = np.sqrt(np.sum(grid**2, axis=0)).reshape([1, height, width]) # avoid zero norm
def main(): train_set = GazeDataset(root_dir='../GazeFollowData/', mat_file='../GazeFollowData/train_annotations.mat', training='train') train_data_loader = DataLoader(train_set, batch_size=32 * 4, shuffle=True, num_workers=16) test_set = GazeDataset(root_dir='../GazeFollowData/', mat_file='../GazeFollowData/test2_annotations.mat', training='test') test_data_loader = DataLoader(test_set, batch_size=32 * 4, shuffle=False, num_workers=8) net = GazeNet() net = DataParallel(net) net.cuda() resume_training = False if resume_training: pretrained_dict = torch.load('../model/pretrained_model.pkl') model_dict = net.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } model_dict.update(pretrained_dict) net.load_state_dict(model_dict) test(net, test_data_loader) exit() method = 'Adam' learning_rate = 0.0001 optimizer_s1 = optim.Adam( [{ 'params': net.module.face_net.parameters(), 'initial_lr': learning_rate }, { 'params': net.module.face_process.parameters(), 'initial_lr': learning_rate }, { 'params': net.module.eye_position_transform.parameters(), 'initial_lr': learning_rate }, { 'params': net.module.fusion.parameters(), 'initial_lr': learning_rate }], lr=learning_rate, weight_decay=0.0001) optimizer_s2 = optim.Adam([{ 'params': net.module.fpn_net.parameters(), 'initial_lr': learning_rate }], lr=learning_rate, weight_decay=0.0001) optimizer_s3 = optim.Adam([{ 'params': net.parameters(), 'initial_lr': learning_rate }], lr=learning_rate * 0.1, weight_decay=0.0001) lr_scheduler_s1 = optim.lr_scheduler.StepLR(optimizer_s1, step_size=5, gamma=0.1, last_epoch=-1) lr_scheduler_s2 = optim.lr_scheduler.StepLR(optimizer_s2, step_size=5, gamma=0.1, last_epoch=-1) lr_scheduler_s3 = optim.lr_scheduler.StepLR(optimizer_s3, step_size=5, gamma=0.1, last_epoch=-1) max_epoch = 25 epoch = 0 while epoch < max_epoch: if epoch == 0: lr_scheduler = lr_scheduler_s1 optimizer = optimizer_s1 elif epoch == 7: lr_scheduler = lr_scheduler_s2 optimizer = optimizer_s2 elif epoch == 15: lr_scheduler = lr_scheduler_s3 optimizer = optimizer_s3 lr_scheduler.step() running_loss = [] for i, data in tqdm(enumerate(train_data_loader)): image, face_image, gaze_field, eye_position, gt_position, gt_heatmap = \ data['image'], data['face_image'], data['gaze_field'], data['eye_position'], data['gt_position'], data['gt_heatmap'] image, face_image, gaze_field, eye_position, gt_position, gt_heatmap = \ map(lambda x: Variable(x.cuda()), [image, face_image, gaze_field, eye_position, gt_position, gt_heatmap]) #for var in [image, face_image, gaze_field, eye_position, gt_position]: # print var.shape optimizer.zero_grad() direction, predict_heatmap = net( [image, face_image, gaze_field, eye_position]) heatmap_loss, m_angle_loss = \ F_loss(direction, predict_heatmap, eye_position, gt_position, gt_heatmap) if epoch == 0: loss = m_angle_loss elif epoch >= 7 and epoch <= 14: loss = heatmap_loss else: loss = m_angle_loss + heatmap_loss loss.backward() optimizer.step() running_loss.append( [heatmap_loss.data[0], m_angle_loss.data[0], loss.data[0]]) if i % 10 == 9: logging.info('%s %s %s' % (str(np.mean(running_loss, axis=0)), method, str(lr_scheduler.get_lr()))) running_loss = [] epoch += 1 save_path = '../model/two_stage_fpn_concat_multi_scale_' + method if not os.path.exists(save_path): os.makedirs(save_path) torch.save(net.state_dict(), save_path + '/model_epoch{}.pkl'.format(epoch)) test(net, test_data_loader)
def main(): dataset_path = '/home/emannuell/Documentos/mestrado/dataset/data_new/' output_path = 'output/convergeTest' train_set = GazeDataset(root_dir=dataset_path, training='train') train_data_loader = DataLoader(train_set, batch_size=10, shuffle=True, num_workers=4) test_set = GazeDataset(root_dir=dataset_path, training='test') test_data_loader = DataLoader(test_set, batch_size=2, shuffle=False, num_workers=4) net = GazeNet() net = DataParallel(net) net.cuda() resume_training = False if resume_training: pretrained_dict = torch.load('model/pretrained_model.pkl') model_dict = net.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } model_dict.update(pretrained_dict) net.load_state_dict(model_dict) test(net, test_data_loader) exit() method = 'Adam' # 0.0001 learning_rate = 0.001 optimizer_s1 = optim.Adam( [{ 'params': net.module.head_pose_transform.parameters(), 'initial_lr': learning_rate }, { 'params': net.module.eye_position_transform.parameters(), 'initial_lr': learning_rate }, { 'params': net.module.fusion.parameters(), 'initial_lr': learning_rate }], lr=learning_rate, weight_decay=0.0001) optimizer_s2 = optim.Adam([{ 'params': net.module.fpn_net.parameters(), 'initial_lr': learning_rate }], lr=learning_rate, weight_decay=0.0001) optimizer_s3 = optim.Adam([{ 'params': net.parameters(), 'initial_lr': learning_rate }], lr=learning_rate * 0.1, weight_decay=0.0001) lr_scheduler_s1 = optim.lr_scheduler.StepLR(optimizer_s1, step_size=5, gamma=0.1, last_epoch=-1) lr_scheduler_s2 = optim.lr_scheduler.StepLR(optimizer_s2, step_size=5, gamma=0.1, last_epoch=-1) lr_scheduler_s3 = optim.lr_scheduler.StepLR(optimizer_s3, step_size=5, gamma=0.1, last_epoch=-1) max_epoch = 25 epoch = 0 while epoch < max_epoch: if epoch == 0: lr_scheduler = lr_scheduler_s1 optimizer = optimizer_s1 # 5 elif epoch == 15: lr_scheduler = lr_scheduler_s2 optimizer = optimizer_s2 # 9 elif epoch == 20: lr_scheduler = lr_scheduler_s3 optimizer = optimizer_s3 # optimizer.step() lr_scheduler.step() running_loss = [] ep_heatmap_loss = [] ep_m_angle_loss = [] for i, data in tqdm(enumerate(train_data_loader)): image, gaze_field, eye_position, gt_position, gt_heatmap, head_pose = \ data['image'], data['gaze_field'], data['eye_position'], data['gt_position'], data['gt_heatmap'], data['head_pose'] image, gaze_field, eye_position, gt_position, gt_heatmap, head_pose = \ map(lambda x: Variable(x.cuda()), [image, gaze_field, eye_position, gt_position, gt_heatmap, head_pose]) optimizer.zero_grad() direction, predict_heatmap = net( [image, gaze_field, eye_position, head_pose]) heatmap_loss, m_angle_loss = F_loss(direction, predict_heatmap, eye_position, gt_position, gt_heatmap) ep_heatmap_loss.append(np.array(heatmap_loss.cpu().data)) ep_m_angle_loss.append(np.array(m_angle_loss.cpu().data)) if epoch == 0: loss = m_angle_loss elif epoch >= 15 and epoch <= 20: loss = heatmap_loss else: loss = m_angle_loss + heatmap_loss loss.backward() optimizer.step() running_loss.append( [heatmap_loss.data, m_angle_loss.data, loss.data]) # if i % 10 == 9: # logging.info('%s %s %s'%(str(np.mean(running_loss, axis=0)), method, str(lr_scheduler.get_lr()))) # running_loss = [] epoch += 1 print('==== Training loss ====') logging.info('Epoch: %s' % epoch) logging.info('heatmap loss: %s' % str(np.mean(np.array(ep_heatmap_loss)))) logging.info('mean angle loss: %s' % str(np.mean(np.array(ep_m_angle_loss)))) logging.info('file: %s' % output_path + '/epoch_{}_loss_{}.pkl'.format(epoch, loss.data)) if not os.path.exists(output_path): os.makedirs(output_path) print('Saving model to output path: ', output_path + '/epoch_{}_loss_{}.pkl'.format(epoch, loss.data)) torch.save( net.state_dict(), output_path + '/epoch_{}_loss_{}.pkl'.format(epoch, loss.data))
def deepMain(): # from deepface.detectors.detector_ssd import FaceDetectorSSDMobilenetV2 net = GazeNet() net = DataParallel(net) net.cuda() # face_detector = FaceDetectorSSDMobilenetV2() face_detector = cv2.CascadeClassifier( 'model/lbpcascade_frontalface_improved.xml') # Load pretrained gaze following model pretrained_dict = torch.load('model/epoch_15_loss_0.0558342523873.pkl') model_dict = net.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } model_dict.update(pretrained_dict) net.load_state_dict(model_dict) cap = cv2.VideoCapture('video.mp4') width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) videoFrames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) fps = int(cap.get(cv2.CAP_PROP_FPS)) print(width, height, videoFrames, fps) # Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file. # out = cv2.VideoWriter('inspecaoCarroResult.avi',cv2.VideoWriter_fourcc('M','J','P','G'), fps, (width, height)) print('Iniciando processamento do video...') while True: frameId = int(cap.get(cv2.CAP_PROP_POS_FRAMES)) if frameId < videoFrames: print(frameId, '/', videoFrames) ret, img = cap.read() t = time.time() originalImg = img img = cv2.resize(img, (640, 480)) height, width, _ = img.shape faces = face_detector.detectMultiScale(img, 1.05, 3) for face in faces: print(face) # Precisa redimensionar imagem para aumentar o crop! faceImage = img[face[1]:face[1] + face[2], face[0]:face[3] + face[0]] gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) formatDetected = [(face[1], face[1] + face[2]), (face[0], face[3] + face[0])] print(formatDetected) x_ = face[0] y_ = face[1] w_ = face[2] h_ = face[3] score = face.score x, y = int(x_ + w_ / 2), int(y_ + h_ / 2) print(x, y) cv2.rectangle(img, (x_, y_), (x_ + w_, y_ + h_), (0, 0, 255), 2) # Head pose estimation faceBbox = x_, y_, x_ + w_, y_ + h_ y, p, r = headpose.detectHeadPose(img, faceBbox) # To normalize detections we fit with min and max angular examples from dataset MinMaxY = [-86.73414612, 87.62715149] MinMaxP = [-54.0966301, 36.50032043] MinMaxR = [-42.67565918, 42.16217041] MinMaxY.append(y) MinMaxP.append(p) MinMaxR.append(r) y = min_max_scaler.fit_transform( np.array(MinMaxY).reshape((-1, 1))) p = min_max_scaler.fit_transform( np.array(MinMaxP).reshape((-1, 1))) r = min_max_scaler.fit_transform( np.array(MinMaxR).reshape((-1, 1))) headPoseAngles = float(y[-1]), float(p[-1]), float(r[-1]) print('Head pose normalized: ', headPoseAngles) center_x = eye_center[0] / width center_y = eye_center[1] / height img = cv2.circle(img, (x, y), 2, (255, 0, 255), thickness=2) # Magento heatmap, p_x, p_y = test(net, originalImg, (x, y), headPoseAngles) img = cv2.circle(img, (int(p_x * width), int(p_y * height)), 2, (255, 0, 0), thickness=2) # Azul img = draw_result(img, (center_x, center_y), heatmap, (p_x, p_y)) else: break img = np.concatenate((img, img2), axis=1) img2 = img # Write the frame into the file 'output.avi' # img = cv2.resize(img, (1280, 480)) # out.write(img) cv2.imshow('result', img) if cv2.waitKey(1) & 0xFF == ord('q'): break cap.release() # out.release() cv2.destroyAllWindows()
def main(): # DEFINE PARAMETERS CHECKPOINT_PATH = '/media/samsung2080pc/New Volume/SAMSUNG/gazefollowing/trial1_Adam' EPOCH = 25 LOAD_PATH = os.path.join(CHECKPOINT_PATH, 'model_epoch' + str(EPOCH) + '.pkl') DATASET_PATH = '/home/samsung2080pc/Documents/ObjectOfInterestV22Dataset' TEST_PATH = '/home/samsung2080pc/Documents/ObjectOfInterestV22Dataset/test.pickle' NUM_TEST = 0 net = GazeNet() net = DataParallel(net) net.cuda() # pretrained_dict = torch.load('../model/pretrained_model.pkl') pretrained_dict = torch.load(LOAD_PATH) model_dict = net.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } model_dict.update(pretrained_dict) net.load_state_dict(model_dict) f = open(TEST_PATH, 'rb') test_data = pickle.load(f) test_img_path = test_data[0]['filename'] test_img_path = os.path.join(DATASET_PATH, test_img_path) h, w = cv2.imread(test_img_path).shape[:2] if NUM_TEST == 0: NUM_TEST = len(test_data) save_path = os.path.join(CHECKPOINT_PATH, 'epoch_' + str(EPOCH)) if not os.path.exists(save_path): os.makedirs(save_path) start_time = time.time() results_pkl = [] for i in tqdm(range(NUM_TEST)): save_path = os.path.join(CHECKPOINT_PATH, 'epoch_' + str(EPOCH), 'out' + str(i).zfill(4) + '.png') # test_image_path = sys.argv[1] # x = float(sys.argv[2]) # y = float(sys.argv[3]) test_image_path = test_data[i]['filename'] test_image_path = os.path.join(DATASET_PATH, test_image_path) x = test_data[i]['hx'] / w y = test_data[i]['hy'] / h # print(test_image_path,x,y) heatmap, p_x, p_y = test(net, test_image_path, (x, y)) output = { # PREDICTIONS 'predictions': { 'heatmap': heatmap, 'p_x': p_x, 'p_y': p_y }, # INPUTS 'inputs': { 'image_path': test_data[i]['filename'], 'eye_x': x, 'eye_y': y, }, # GROUND TRUTH 'gt': { 'gaze_cx': test_data[i]['gaze_cx'], 'gaze_cy': test_data[i]['gaze_cy'] } } results_pkl.append(output) draw_result(test_image_path, (x, y), heatmap, (p_x, p_y), save_path) end_time = time.time() process_time = end_time - start_time outfilename = os.path.join(CHECKPOINT_PATH, 'epoch_' + str(EPOCH), 'allresults.pkl') with open(outfilename, 'wb') as outfile: pickle.dump(results_pkl, outfile) print(results_pkl) print('Processed %i images in %f seconds.' % (NUM_TEST, process_time))
def main(): ''' train_set = GazeDataset(root_dir='../../data/', mat_file='../../data/train_annotations.mat', training='train') train_data_loader = DataLoader(train_set, batch_size=48, shuffle=True, num_workers=8) test_set = GazeDataset(root_dir='../../test_data/', mat_file='../../test_data/test2_annotations.mat', training='test') test_data_loader = DataLoader(test_set, batch_size=32, shuffle=False, num_workers=8) ''' dis_train_sets = dataset_wrapper( root_dir='../../data/', mat_file='../../data/train_annotations.mat', training='train') #dis_train_data_loader = DataLoader(dis_train_sets[0], batch_size=48, # shuffle=True, num_workers=8) dis_test_sets = dataset_wrapper( root_dir='../../test_data/', mat_file='../../test_data/test2_annotations.mat', training='test') #dis_test_data_loader = DataLoader(dis_test_sets[0], batch_size=32, # shuffle=False, num_workers=8) dis_train_data_loaders, dis_test_data_loaders = [], [] for i in range(16): dis_train_data_loaders.append( DataLoader(dis_train_sets[i], batch_size=40, shuffle=True, num_workers=8)) dis_test_data_loaders.append( DataLoader(dis_test_sets[i], batch_size=16, shuffle=False, num_workers=1)) net = GazeNet() net = DataParallel(net) net.cuda() #print(next(net.module.fpn_net.parameters()).is_cuda) ##print(next(net.module.fpn_net.parameters()).is_cuda) area_count = 8 area_in_network = int(16 / area_count) cur_area_idx = 0 fpn_weights_transferred = False for i in range(area_count): net.module.change_fpn(i) if not next(net.module.fpn_net.parameters()).is_cuda: net.module.fpn_net.cuda() net.module.change_fpn(cur_area_idx) ##print(next(net.module.fpn_net.parameters()).is_cuda) #exit(0) resume_training = False if resume_training: pretrained_dict = torch.load('../model/pretrained_model.pkl') model_dict = net.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } model_dict.update(pretrained_dict) net.load_state_dict(model_dict) test(net, test_data_loader) exit() method = 'Adam' learning_rate = 0.0001 optimizer_s1 = optim.Adam( [{ 'params': net.module.face_net.parameters(), 'initial_lr': learning_rate }, { 'params': net.module.face_process.parameters(), 'initial_lr': learning_rate }, { 'params': net.module.eye_position_transform.parameters(), 'initial_lr': learning_rate }, { 'params': net.module.fusion.parameters(), 'initial_lr': learning_rate }], lr=learning_rate, weight_decay=0.0001) #optimizer_s2 = optim.Adam([{'params': net.module.fpn_net.parameters(), # 'initial_lr': learning_rate}], # lr=learning_rate, weight_decay=0.0001) optimizer_s2s, optimizer_s3s = [], [] for i in range(area_count): net.module.change_fpn(i) optimizer_s2 = optim.Adam( [{ 'params': net.module.fpn_nets[i].parameters(), 'initial_lr': learning_rate }], lr=learning_rate, weight_decay=0.0001) optimizer_s3 = optim.Adam([{ 'params': net.parameters(), 'initial_lr': learning_rate }], lr=learning_rate * 0.1, weight_decay=0.0001) optimizer_s2s.append(optimizer_s2) optimizer_s3s.append(optimizer_s3) optimizer_s2 = optimizer_s2s[0] optimizer_s3 = optimizer_s3s[0] lr_scheduler_s1 = optim.lr_scheduler.StepLR(optimizer_s1, step_size=5, gamma=0.1, last_epoch=-1) #lr_scheduler_s2 = optim.lr_scheduler.StepLR(optimizer_s2, step_size=5, gamma=0.1, last_epoch=-1) lr_scheduler_s2s, lr_scheduler_s3s = [], [] for i in range(area_count): lr_scheduler_s2 = optim.lr_scheduler.StepLR(optimizer_s2s[i], step_size=5, gamma=0.1, last_epoch=-1) lr_scheduler_s3 = optim.lr_scheduler.StepLR(optimizer_s3s[i], step_size=5, gamma=0.1, last_epoch=-1) lr_scheduler_s2s.append(lr_scheduler_s2) lr_scheduler_s3s.append(lr_scheduler_s3) lr_scheduler_s2 = lr_scheduler_s2s[0] lr_scheduler_s3 = lr_scheduler_s3s[0] # Set the model to use the first FPN net.module.change_fpn(cur_area_idx) max_epoch = 30 epoch = 0 #epoch = 7 while epoch < max_epoch: logging.info('\n--- Epoch: %s\n' % str(epoch)) if epoch == 0: lr_scheduler = lr_scheduler_s1 optimizer = optimizer_s1 elif epoch == 7: lr_scheduler = lr_scheduler_s2 optimizer = optimizer_s2 elif epoch == 15: lr_scheduler = lr_scheduler_s3 optimizer = optimizer_s3 #lr_scheduler.step() #lr_scheduler.step() running_loss = [] #for data_loader_idx in range(len(dis_train_data_loaders)): for data_loader_idx in range(len(dis_train_data_loaders)): train_data_loader = dis_train_data_loaders[data_loader_idx] if epoch >= 10: #if epoch >= 7: if not fpn_weights_transferred: net.module.transfer_fpn_weights() fpn_weights_transferred = True area_idx = int(data_loader_idx / area_in_network) if cur_area_idx != area_idx: cur_area_idx = area_idx net.module.change_fpn(cur_area_idx) if epoch < 15: lr_scheduler = lr_scheduler_s2s[cur_area_idx] optimizer = optimizer_s2s[cur_area_idx] else: lr_scheduler = lr_scheduler_s3s[cur_area_idx] optimizer = optimizer_s3s[cur_area_idx] #if not next(net.module.fpn_net.parameters()).is_cuda: # net.module.fpn_net.cuda() #test_data_loader = dis_test_data_loaders[data_loader_idx] #train_data_loader = DataLoader(dis_train_sets[data_loader_idx], batch_size=48, # shuffle=True, num_workers=2) #test_data_loaders = DataLoader(dis_test_sets[data_loader_idx], batch_size=32, # shuffle=False, num_workers=2) for i, data in tqdm(enumerate(train_data_loader)): image, face_image, gaze_field, eye_position, gt_position, gt_heatmap = \ data['image'], data['face_image'], data['gaze_field'], data['eye_position'], data['gt_position'], data['gt_heatmap'] image, face_image, gaze_field, eye_position, gt_position, gt_heatmap = \ map(lambda x: x.cuda(), [image, face_image, gaze_field, eye_position, gt_position, gt_heatmap]) # for var in [image, face_image, gaze_field, eye_position, gt_position]: # print var.shape optimizer.zero_grad() direction, predict_heatmap = net( [image, face_image, gaze_field, eye_position]) heatmap_loss, m_angle_loss = \ F_loss(direction, predict_heatmap, eye_position, gt_position, gt_heatmap) if epoch == 0: #if epoch < 7: loss = m_angle_loss elif epoch >= 7 and epoch <= 14: loss = heatmap_loss else: loss = m_angle_loss + heatmap_loss loss.backward() optimizer.step() # running_loss.append([heatmap_loss.data[0], # m_angle_loss.data[0], loss.data[0]]) running_loss.append( [heatmap_loss.item(), m_angle_loss.item(), loss.item()]) if i % 10 == 9: logging.info('%s %s %s' % (str(np.mean(running_loss, axis=0)), method, str(lr_scheduler.get_last_lr()))) running_loss = [] lr_scheduler.step() epoch += 1 save_path = '../model/two_stage_fpn_concat_multi_scale_' + method if not os.path.exists(save_path): os.makedirs(save_path) if epoch % 5 == 0: torch.save(net.state_dict(), save_path + '/model_epoch{}.pkl'.format(epoch)) for i in range(16): torch.save(net.module.fpn_nets[i].state_dict(), save_path + '/fpn_{}.pkl'.format(i)) for data_loader_idx in range(len(dis_test_data_loaders)): test_data_loader = dis_test_data_loaders[data_loader_idx] if epoch > 10: area_idx = int(data_loader_idx / area_in_network) net.module.change_fpn(area_idx) cur_area_idx = area_idx test(net, test_data_loader)