In [ ]:

    

import os, glob, platform, datetime, random
from collections import OrderedDict

import torch
import torch.nn as nn
import torch.utils.data as data_utils
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
from torch.autograd import Variable
from torch import functional as F
# import torchvision.datasets as datasets
import torchvision.models as models
import torchvision.transforms as transforms
import scipy.io as sio
import cv2
from PIL import Image
from tensorboardX import SummaryWriter

import numpy as np
from numpy.linalg import inv as denseinv
from scipy import sparse
from scipy.sparse import lil_matrix, csr_matrix
from scipy.sparse.linalg import spsolve
from scipy.sparse.linalg import inv as spinv
import scipy.misc

from myimagefoldereccv import MyImageFolder
from mymodel import GradientNet
from myargs import Args


    

In [ ]:

    

def loadimg(path):
    im = Image.open(path).convert('RGB')
    print(im.size)
    im = transforms.ToTensor()(im)
    x = torch.zeros(1,3,416,1024)
    x[0,:,:,:] = im[:,0:416,0:1024]
    #x = torch.zeros(1,3,32,32)
    #x[0,:,:,:] = im[:,0:32,0:32]
    x = Variable(x, volatile=True)
    return x


    

In [ ]:

    

def save_csv(path, para):
    text = ''
    n,c,h,w = para.size()
    text += ','.join([str(n), str(c), str(h), str(w)]) + ','
    for nn in range(n):
        for cc in range(c):
            for hh in range(h):
                for ww in range(w):
                    text += str(para[nn,cc,hh,ww].data.cpu().numpy()) + ','
    with open(path, 'w') as f:
        f.write(text)


    

In [ ]:

    

gpu_num = 0
gradient = False
type2 = 'rgb' if gradient == False else 'gd'
image_slpit = True

# parameters = filter(lambda p: p.requires_grad, net.parameters())
# optimizer = optim.SGD(parameters, lr=args.base_lr, momentum=args.momentum)

res_root = './results/images/'
scenes = glob.glob('/home/albertxavier/dataset/sintel2/clean/*')
cnt_albedo = 0
cnt_shading = 0
for scene in scenes:
    scene = scene.split('/')[-1]
    res_dir = os.path.join(res_root, 'image_split', scene)
    if not os.path.exists(res_dir):
        os.makedirs(res_dir)
    for type_ in ['albedo', 'shading']:
        
#         if scene!='image_split': continue
        
        #root = '/media/lwp/xavier/graduation_results/showcase_model/image_split/{}/{}/'.format(type_, type2)
#         root = '/media/lwp/xavier/graduation_results/showcase_model/{}/{}/{}/'.format(scene, type_, type2)
        root = '/media/albertxavier/data/eccv/graduation-project/pytorch/snapshot0/'
        print (root+'snapshot-238.pth.tar')
        if not os.path.exists(root+'snapshot-238.pth.tar'): continue
        snapshot = torch.load(root+'snapshot-238.pth.tar')
        state_dict = snapshot['state_dict']
        args = snapshot['args']
        densenet = models.__dict__[args.arch](pretrained=True).cuda(gpu_num)
        
#         net.load_state_dict(state_dict)
#         net.train()
        net = None
        num = 40 if scene=='market_6' else 50
        for ind in range(1, 11):
            if net is not None: del net
#             torch.cuda.empty_cache()
            net = GradientNet(densenet=densenet, growth_rate=32, 
                          transition_scale=2, pretrained_scale=4,
                    debug=False).cuda(gpu_num)
            
            net.load_state_dict(state_dict)
            net.train()
            frame = 'frame_%04d.png'%(ind)
            print('/home/albertxavier/dataset/sintel2/clean/{}/{}'.format(scene, frame))
            im = loadimg('/home/albertxavier/dataset/sintel2/clean/{}/{}'.format(scene, frame)).cuda(gpu_num)
            print(im.size())
            merged = net(im.cuda(gpu_num))
            alpha = merged[:,9:10,:,:]
            beta = merged[:,10:13,:,:]
            print('alpha', alpha.min(), alpha.max())
            print('beta', beta.min(), beta.max())
            #print(merged)
            ######
            #break
            ######
            # merged = mergeRGB
            merged = merged[0]
            merged = merged[0:3,:,:]
            merged = merged.cpu().data.numpy()
            print (merged.shape)
            merged = merged.transpose(1,2,0)
            print (merged.shape)
            B = merged[:,:,0:3]
            dx = merged[:,:,3:6]
            dy = merged[:,:,6:9]
            res_frame = 'albedo_%04d.png'%(ind) if type_ == 'albedo' else 'shading_%04d.png'%(ind)
            res_dx_frame = 'albedo_dx_%04d.png'%(ind) if type_ == 'albedo' else 'shading_%04d.png'%(ind)
            res_dy_frame = 'albedo_dy_%04d.png'%(ind) if type_ == 'albedo' else 'shading_%04d.png'%(ind)
            res_alpha = 'alpha_%04d.mat'%(ind)
            res_beta =  'beta_%04d.mat'%(ind)
            print('res path', os.path.join(res_dir,res_frame))
            cv2.imwrite(os.path.join(res_dir,res_frame), B[:,:,::-1]*255)   
            cv2.imwrite(os.path.join(res_dir,res_dx_frame), dx[:,:,::-1]*255)   
            cv2.imwrite(os.path.join(res_dir,res_dy_frame), dy[:,:,::-1]*255)   
            # save_csv(os.path.join(res_dir,res_alpha), alpha)
            # save_csv(os.path.join(res_dir,res_beta), beta)
            sio.savemat(os.path.join(res_dir,res_alpha), {'alpha': alpha.cpu().data.numpy()})
            sio.savemat(os.path.join(res_dir,res_beta), {'beta': beta.cpu().data.numpy()})
            break


    

In [ ]:

    

# if gradient == False:
#     merged = mergeRGB[5]
#     merged = merged[0]
#     merged = merged.cpu().data.numpy()
#     print (merged.shape)
#     merged = merged.transpose(1,2,0)
#     print (merged.shape)
#     dx = merged[:,:,0:3]
#     cv2.imwrite('out_merge.png', dx[:,:,::-1]*255)


    

In [ ]:

    

# if gradient == True:
#     merged = mergeRGB[5]
#     merged = merged[0]
#     merged = merged.cpu().data.numpy()
#     print (merged.shape)
#     merged = merged.transpose(1,2,0)
#     print (merged.shape)
#     dy = merged[:,:,0:3]+0.5
#     dx = merged[:,:,3:6]+0.5
#     cv2.imwrite('out_merge_dx.png', dx[:,:,::-1]*255)
#     cv2.imwrite('out_merge_dy.png', dy[:,:,::-1]*255)