def main():
model = define_model(is_resnet=False, is_densenet=False, is_senet=True)
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(torch.load('./pretrained_model/model_senet'))
model.eval()
DetList = json.load(open('data/test/hoi_01.json'))
indir = 'data/test/hoi_01'
inframes = [
f for f in os.listdir(indir) if os.path.isfile(os.path.join(indir, f))
]
inframes.sort(key=lambda x: int(x.split('.')[0]))
for frame in inframes:
fid = frame.split('.')[0]
fname = os.path.join(indir, frame)
ori_frame = Image.open(fname).convert("RGB")
if str(int(fid)) not in DetList.keys():
outframe = ori_frame
outframe.save('data/test/outputs/' + fid + '.jpg')
print(fid)
else:
# ori_size = ori_frame.size
nyu2_loader = loaddata.readNyu2(fname)
detframe = DetList[str(int(fid))]
with torch.no_grad():
test(fid, detframe, nyu2_loader, model, ori_frame)
def main():
# choose the backbone architecture:
is_resnet = False
is_densenet = True
is_senet = False
model = define_model(is_resnet, is_densenet, is_senet)
model = torch.nn.DataParallel(model).cuda()
# load the pretrained model:
if is_resnet:
checkpoint = torch.load('./pretrained_model/resnet50_checkpoint.pth.tar')
# model.load_state_dict(torch.load('./pretrained_model/model_resnet'))
if is_densenet:
checkpoint = torch.load('./pretrained_model/densenet_checkpoint.pth.tar')
# model.load_state_dict(torch.load('./pretrained_model/model_densenet'))
if is_senet:
checkpoint = torch.load('./pretrained_model/senet_checkpoint.pth.tar')
# model.load_state_dict(torch.load('./pretrained_model/model_senet'))
model.load_state_dict(checkpoint['state_dict'])
model.eval()
nyu2_loader = loaddata_demo.readNyu2('./data/demo/img_nyu2.png')
# nyu2_loader = loaddata_demo.readNyu2('./data/demo/images')
# print('-----++++transformed image:', type(nyu2_loader))
test(nyu2_loader, model)
def main():
model = define_model(is_resnet=False, is_densenet=False, is_senet=True)
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(torch.load('./pretrained_model/model_senet'))
model.eval()
nyu2_loader = loaddata.readNyu2('data/demo/img_nyu2.png')
test(nyu2_loader, model)
def main():
model = define_model(is_resnet=False, is_densenet=False, is_senet=True)
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(torch.load('epoch-5.pth.tar')['state_dict'])
model.eval()
mypath='data/demo/'
files = [f for f in listdir(mypath) if isfile(join(mypath, f))]
for f in files:
nyu2_loader = loaddata.readNyu2(join(mypath, f))
test(nyu2_loader, model, f)
def main(image_path):
model = define_model(is_resnet=False, is_densenet=False, is_senet=True)
model = torch.nn.DataParallel(model)
if torch.cuda.is_available():
model = model.cuda()
model.load_state_dict(
torch.load(f='./pretrained_model/model_senet',
map_location=None if torch.cuda.is_available() else 'cpu'))
model.eval()
nyu2_loader = loaddata.readNyu2(image_path)
test(nyu2_loader, model)
def main():
if (not os.path.exists(args.output)):
print("Output directory doesn't exist! Creating...")
os.makedirs(args.output)
model = define_model(is_resnet=False, is_densenet=False, is_senet=True)
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(torch.load('./pretrained_model/model_senet'))
model.eval()
print
img = cv2.imread(args.img)
nyu2_loader = loaddata.readNyu2(args.img)
test(nyu2_loader, model, img.shape[1], img.shape[0])
def main():
model = define_model(is_resnet=False, is_densenet=False, is_senet=True)
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(torch.load('./pretrained_model/model_senet'))
model.eval()
image_dir = "/home/taylor/Mirror-Segmentation/data_640/train/image/"
mask_dir = "/home/taylor/Mirror-Segmentation/data_640/train/mask/"
imglist = os.listdir(image_dir)
output_dir = "/media/taylor/mhy/depth_original/train/"
for i, imgname in enumerate(imglist):
print i
mask_path = mask_dir + imgname[:-4] + "_json/label8.png"
output_path = output_dir + imgname[:-4] + ".png"
nyu2_loader = loaddata.readNyu2(os.path.join(image_dir, imgname))
test(nyu2_loader, model, mask_path, output_path)
def test(model, rgb_path, output_path):
nyu2_loader = loaddata.readNyu2(rgb_path)
path, file = os.path.split(rgb_path)
file = f"{file.split('.')[0]}.png"
depth_path = os.path.join(output_path,
file) if output_path else os.path.join(
path, f"out_{file}")
print(f"{rgb_path} -> {depth_path}")
for i, image in enumerate(nyu2_loader):
image = image.cuda()
out = model(image)
matplotlib.image.imsave(
depth_path,
out.view(out.size(2), out.size(3)).data.cpu().numpy())
def main():
imggg = cv2.imread(args.Photo_string)
imggg = cv2.resize(imggg,(my_width,my_height))
#Depth_lever
steps = 10
if len(args.Depth_lever) == 1:
steps = ord(args.Depth_lever[0]) - ord('0')
elif len(args.Depth_lever) == 2:
steps = (ord(args.Depth_lever[0]) - ord('0'))*10 + ord(args.Depth_lever[1]) - ord('0')
print(steps)
model_name = 'senet'
if model_name == 'resnet':
model = define_model(is_resnet=True, is_densenet=False, is_senet=False)
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(torch.load(pm.mum_path+'/pretrained_model/model_resnet'))
model.eval()
elif model_name == 'densenet':
model = define_model(is_resnet=False, is_densenet=True, is_senet=False)
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(torch.load(pm.mum_path+'/pretrained_model/model_densenet'))
model.eval()
elif model_name == 'senet':
model = define_model(is_resnet=False, is_densenet=False, is_senet=True)
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(torch.load(pm.mum_path+'/pretrained_model/model_senet'))
model.eval()
print ("model loaded!")
nyu2_loader = loaddata.readNyu2(imggg)
for i,image in enumerate(nyu2_loader):
with torch.no_grad(): #释放内存
image = torch.autograd.Variable(image, volatile=True).cuda()
out = model(image)
mid_img = out.view(out.size(2),out.size(3)).data.cpu().numpy()
# plt.imshow(mid_img)
# plt.show()
#outt = np.array(mid_img)/np.max(mid_img) - np.min(mid_img)/np.max(mid_img)
outt = (np.array(mid_img) - np.min(mid_img))/(np.max(mid_img) - np.min(mid_img))
#print outt
outt = cv2.resize(outt,(imggg.shape[1],imggg.shape[0]))
#print outt.shape,outt
#cv2.imshow("depth",outt)
# save_img = cv2.resize(imggg,(imggg.shape[1]*2,imggg.shape[0]))
# save_d = cv2.cvtColor(outt*255,cv2.COLOR_GRAY2BGR)
# save_img[:,0:imggg.shape[1]] = imggg.copy()
# save_img[:,imggg.shape[1]:save_img.shape[1]] = save_d.copy()
# cv2.imwrite("test_rgb-d/senet"+args.Photo_string+".jpg",save_img) ##保存深度图
cv2.imwrite("test_rgb-d/d-senet"+args.Photo_string+".jpg",outt*255) ##保存深度图
# # #dibr处理
# #outtt = cv2.imread("test_rgb-d/d-senet"+args.Photo_string+".jpg",0)
# frame = cv2.imread(args.Photo_string)
# frame = cv2.resize(frame,(my_width,my_height))
# imgdd, steps = dibr.depth_gradation(outt*255)
# cv2.imshow("d",imgdd/255)
# left,right,sbs_img = dibr.segementation(frame,imgdd,steps)
# rb_3d = dibr.Merge3D(left,right)
# #cv2.imshow("0",frame)
# cv2.imshow("left",left)
# cv2.imshow("right",right)
# cv2.imshow("d",outt)
# cv2.imwrite(args.Photo_string+"dpth.jpg",outt*255)
cv2.waitKey(0)
cv2.destroyAllWindows()
def start_run():
model_name = args.model
input_img = cv2.imread(args.in_photo)
my_width = input_img.shape[1]
my_height = input_img.shape[0]
radio = 1.0
if my_width > my_height and my_width > 640:
radio = 640.0 / my_width
my_width = 640
my_height = int(radio * my_height)
elif my_width < my_height and my_height > 640:
radio = 640.0 / my_height
my_height = 640
my_width = int(radio * my_width)
#检测 原图输入
radio_rects = []
list_rects = faces_detect(input_img)
print('faces:' + str(len(list_rects)), list_rects)
for (x1, y1, x2, y2) in list_rects:
x1 = int(radio * x1)
y1 = int(radio * y1)
x2 = int(radio * x2)
y2 = int(radio * y2)
radio_rects.append((x1, y1, x2, y2))
print('radio_faces:' + str(len(radio_rects)), radio_rects)
print(my_width, my_height), model_name
input_img = cv2.resize(input_img, (my_width, my_height))
img_depth = cv2.cvtColor(input_img, cv2.COLOR_BGR2GRAY)
if model_name == 'fcrn':
import tensorflow as tf
import NYU_FCRNmodels as models
# Default input size
width = my_width
height = my_height
channels = 3
batch_size = 1
# Create a placeholder for the input image
#placeholder()函数是在神经网络构建graph的时候在模型中的占位，此时并没有把要输入的数据传入模型，它只会分配必要的内存
input_node = tf.placeholder(tf.float32,
shape=(None, height, width, channels))
# Construct the network
net = models.ResNet50UpProj({'data': input_node}, batch_size, 1, False)
#gpu memory set
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.7
with tf.Session(config=config, ) as sess:
starttime = time.time()
# Load the converted parameters
print('Loading the model')
# Use to load from ckpt file
saver = tf.train.Saver()
saver.restore(
sess,
'/home/jiang/Repositories/2.23dchange/depth_image-py/FCRN-DepthPrediction-master/tensorflow/NYU_FCRN.ckpt'
)
print("load time:", time.time() - starttime)
source_img = input_img.copy() #cv2.imread(args.in_photo)
#deal image
img = Image.fromarray(cv2.cvtColor(
source_img, cv2.COLOR_BGR2RGB)) #opencv 转 PIL.Iamge
source_img = cv2.resize(source_img, (my_width, my_height))
img = img.resize([width, height], Image.ANTIALIAS) #缩放
img = np.array(img).astype('float32') #数据类型转换
img = np.expand_dims(np.asarray(img), axis=0) #表示在0位置添加数据
# Evalute the network for the given image
predicttime = time.time()
pred = sess.run(net.get_output(), feed_dict={input_node: img})
img_depth = cv2.cvtColor(np.asarray(pred[0, :, :, 0]),
cv2.COLOR_RGB2BGR)
img_depth = cv2.cvtColor(img_depth, cv2.COLOR_BGR2GRAY)
max_depth = np.max(img_depth)
min_depth = np.min(img_depth)
img_depth = (np.array(img_depth) - min_depth) / (
max_depth - min_depth) #归一化，归为０－１之间
img_depth = cv2.resize(img_depth, (my_width, my_height))
else:
import torch.nn.parallel
from torchvision import transforms, utils
from pytorch_models import modules, net, resnet, densenet, senet
import loaddata_demo as loaddata
import pdb
import model_path as pm
def define_model(is_resnet, is_densenet, is_senet):
if is_resnet:
original_model = resnet.resnet50(pretrained=True)
Encoder = modules.E_resnet(original_model)
model = net.model(Encoder,
num_features=2048,
block_channel=[256, 512, 1024, 2048])
if is_densenet:
original_model = densenet.densenet161(pretrained=True)
Encoder = modules.E_densenet(original_model)
model = net.model(Encoder,
num_features=2208,
block_channel=[192, 384, 1056, 2208])
if is_senet:
original_model = senet.senet154(pretrained='imagenet')
Encoder = modules.E_senet(original_model)
model = net.model(Encoder,
num_features=2048,
block_channel=[256, 512, 1024, 2048])
return model
if 1:
model_name = 'senet'
if model_name == 'resnet':
model = define_model(is_resnet=True,
is_densenet=False,
is_senet=False)
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(
torch.load(pm.mum_path + '/pretrained_model/model_resnet'))
model.eval()
elif model_name == 'densenet':
model = define_model(is_resnet=False,
is_densenet=True,
is_senet=False)
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(
torch.load(pm.mum_path +
'/pretrained_model/model_densenet'))
model.eval()
elif model_name == 'senet':
model = define_model(is_resnet=False,
is_densenet=False,
is_senet=True)
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(
torch.load(pm.mum_path + '/pretrained_model/model_senet'))
model.eval()
print("model loaded!")
nyu2_loader = loaddata.readNyu2(input_img)
for i, image in enumerate(nyu2_loader):
with torch.no_grad(): #释放内存
image = torch.autograd.Variable(image,
volatile=True).cuda()
out = model(image)
mid_img = out.view(out.size(2),
out.size(3)).data.cpu().numpy()
img_depth = (np.array(mid_img) - np.min(mid_img)) / (
np.max(mid_img) - np.min(mid_img))
img_depth = cv2.resize(
img_depth, (input_img.shape[1], input_img.shape[0]))
#cv2.imwrite("faces_result/senet_"+args.in_photo,img_depth*255) ##保存深度图
save_img = input_img.copy()
save_img = cv2.resize(input_img,
(input_img.shape[1] * 2, input_img.shape[0]))
save_d = cv2.cvtColor(img_depth * 255, cv2.COLOR_GRAY2BGR)
#缩放后，乘了比例系数后的rects
mask = cv2.imread("mask1.png")
for (x1, y1, x2, y2) in radio_rects:
cv2.rectangle(save_d, (x1, y1), (x2, y2), (255, 0, 0), 1, cv2.LINE_AA)
pos = (x1, y1)
rect_mean_value = int(np.mean(save_d[y1:y2, x1:x2, 0]))
font = cv2.FONT_HERSHEY_SIMPLEX
cv2.putText(save_d, str(rect_mean_value), pos, font, 0.4, (0, 0, 255),
1, cv2.LINE_AA)
cv2.putText(input_img, str(rect_mean_value), pos, font, 0.4,
(255, 255, 0), 1, cv2.LINE_AA)
mask = cv2.resize(mask, (x2 - x1 - 4, y2 - y1 - 4))
input_img[y1 + 2:y2 - 2, x1 + 2:x2 - 2] = mask.copy()
cv2.rectangle(input_img, (x1, y1), (x2, y2), (0, 0, 255), 1,
cv2.LINE_AA)
save_img[:, 0:input_img.shape[1]] = input_img.copy()
save_img[:, input_img.shape[1]:save_img.shape[1]] = save_d.copy()
cv2.imwrite("faces_depth/" + model_name + '-' + args.in_photo,
save_img) ##保存深度图
cv2.imshow(model_name + "show", save_img)
cv2.waitKey(0)
cv2.destroyAllWindows()
return 0
