def mark_all(model, data_dir, out_dir='marked_predictions', backbone='resnet34',
input_size=None, preprocessing_fcn=None):
if preprocessing_fcn is None:
preprocessing_fcn = get_preprocessing(backbone)
if not os.path.isdir(out_dir):
os.mkdir(out_dir)
images = []
for directory, _, files in os.walk(data_dir):
for fn in files:
if is_image(fn):
images.append(os.path.join(directory, fn))
s = 'filename,x1,y1,x2,y2,x3,y3,x4,y4,x5,y5,x6,y6,x7,y7,x8,y8,x9,y9,x10,y10\n'
for fn_full in images:
name = os.path.splitext(os.path.basename(fn_full))[0]
mask = predict(model=model,
image=fn_full,
out_path=None,
preprocessing_fcn=preprocessing_fcn,
input_size=input_size)
pts = pmask2points(np.squeeze(mask))
s += (os.path.basename(fn_full) + ',' +
','.join('%s,%s' % (x, y) for x, y in pts) + '\n')
image = resize(cv.imread(fn_full), dsize=input_size)
cv.imwrite(os.path.join(out_dir, name + '.jpg'),
mark_image(image, pts))
cv.imwrite(os.path.join(out_dir, name + '.png'),
(np.squeeze(mask).sum(axis=2)*255).astype('uint8'))
with open(os.path.join(out_dir, 'points.csv'), 'w+') as csv_out:
csv_out.write(s)
def rename(folder, format):
"""
This function renames all images in a given
directory according to the pattern given in
the format given, init.
Format:
use NUM for the number of the image.
imgNUM -> img1, img2, .., img100,..
"""
# Get list of all files then rename images
path = util.abspath(folder)
names_lis = os.listdir(path)
os.chdir(path) # cwd to work with OS commands
# Scan for images and rename them
for i in range(0, len(names_lis)):
f = names_lis[i]
extension = get_extension(f)
location = "".join([path, "/", f])
# Need absolute location
#print("checking location %s. Is file? %s " % (location, is_image(location)))
if util.is_image(location):
new_f = "".join([format.replace("NUM", str(i)), ".", extension])
util.update_line("Renaming file %s to %s." % (f, new_f))
os.rename(f, new_f) #Tada, done
def predict_all(model,
data_dir,
out_dir='results',
backbone='resnet34',
input_size=None,
preprocessing_fcn=None):
if preprocessing_fcn is None:
preprocessing_fcn = get_preprocessing(backbone)
if not os.path.isdir(out_dir):
os.mkdir(out_dir)
images = []
for directory, _, files in os.walk(data_dir):
for fn in files:
if is_image(fn):
images.append(os.path.join(directory, fn))
for fn_full in images:
name = os.path.splitext(os.path.basename(fn_full))[0]
predict(model=model,
image=fn_full,
out_path=os.path.join(out_dir, name + '.png'),
preprocessing_fcn=preprocessing_fcn,
input_size=input_size)
def __init__(self, image_dir):
super(DataLoaderHelper, self).__init__()
self.albedo_path = join(image_dir, "albedo")
self.depth_path = join(image_dir, "depth")
self.direct_path = join(image_dir, "direct")
self.normal_path = join(image_dir, "normal")
self.gt_path = join(image_dir, "gt")
self.image_filenames = [
x for x in listdir(self.albedo_path) if is_image(x)
]
def __init__(self, image_dir):
super(DataLoaderHelper, self).__init__()
self.albedo_path = join(image_dir, "albedo")
self.depth_path = join(image_dir, "depth")
self.direct_path = join(image_dir, "direct")
self.normal_path = join(image_dir, "normal")
self.gt_path = join(image_dir, "gt")
self.image_filenames = [x for x in listdir(self.gt_path) if is_image(x)]
for i, im_name in enumerate(self.image_filenames):
base = im_name.split('-')
im_name = base[1]+'-'+base[2]
self.image_filenames[i] = im_name
type=int,
default=3,
help='output channel')
parser.add_argument('--n_generator_filters',
type=int,
default=64,
help="number of generator filters")
opt = parser.parse_args()
netG_model = torch.load(opt.model)
netG = G(opt.n_channel_input * 4, opt.n_channel_output,
opt.n_generator_filters)
netG.load_state_dict(netG_model['state_dict_G'])
root_dir = 'dataset/{}/test/'.format(opt.dataset)
image_dir = 'dataset/{}/test/albedo'.format(opt.dataset)
image_filenames = [x for x in os.listdir(image_dir) if is_image(x)]
for image_name in image_filenames:
albedo_image = load_image(root_dir + 'albedo/' + image_name)
direct_image = load_image(root_dir + 'direct/' + image_name)
normal_image = load_image(root_dir + 'normal/' + image_name)
depth_image = load_image(root_dir + 'depth/' + image_name)
gt_image = load_image(root_dir + 'gt/' + image_name)
albedo = Variable(albedo_image).view(1, -1, 256, 256).cuda()
direct = Variable(direct_image).view(1, -1, 256, 256).cuda()
normal = Variable(normal_image).view(1, -1, 256, 256).cuda()
depth = Variable(depth_image).view(1, -1, 256, 256).cuda()
netG = netG.cuda()