39 lines
894 B
Python
39 lines
894 B
Python
import os
|
|
import torch
|
|
|
|
IMG_EXTENSIONS = [
|
|
'.jpg', '.JPG', '.jpeg', '.JPEG',
|
|
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
|
|
]
|
|
|
|
|
|
def is_image_file(filename):
|
|
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
|
|
|
|
|
|
def make_dataset(dir):
|
|
images = []
|
|
assert os.path.isdir(dir), '%s is not a valid directory' % dir
|
|
|
|
for root, _, fnames in sorted(os.walk(dir)):
|
|
for fname in fnames:
|
|
if is_image_file(fname):
|
|
path = os.path.join(root, fname)
|
|
images.append(path)
|
|
|
|
return images
|
|
|
|
|
|
def edge_compute(x):
|
|
x_diffx = torch.abs(x[:,:,1:] - x[:,:,:-1])
|
|
x_diffy = torch.abs(x[:,1:,:] - x[:,:-1,:])
|
|
|
|
y = x.new(x.size())
|
|
y.fill_(0)
|
|
y[:,:,1:] += x_diffx
|
|
y[:,:,:-1] += x_diffx
|
|
y[:,1:,:] += x_diffy
|
|
y[:,:-1,:] += x_diffy
|
|
y = torch.sum(y,0,keepdim=True)/3
|
|
y /= 4
|
|
return y |