整合去雾网页工具

GCANet/utils.py

@@ -0,0 +1,39 @@
+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