Pārlūkot izejas kodu

support MAPS dataset

Youngwoon Lee 6 gadi atpakaļ
vecāks
revīzija
11a2db91c4
1 mainītis faili ar 8 papildinājumiem un 4 dzēšanām
  1. 8 4
      data_loader.py

+ 8 - 4
data_loader.py

@@ -16,10 +16,14 @@ def read_image(path):
 
     # split image
     h, w, c = image.shape
-    assert w == 256 or w == 512, 'Image size mismatch ({}, {})'.format(h, w)
-    assert h == 128 or h == 256, 'Image size mismatch ({}, {})'.format(h, w)
-    image_a = image[:, :w/2, :].astype(np.float32) / 255.0
-    image_b = image[:, w/2:, :].astype(np.float32) / 255.0
+    assert w in [256, 512, 1200], 'Image size mismatch ({}, {})'.format(h, w)
+    assert h in [128, 256, 600], 'Image size mismatch ({}, {})'.format(h, w)
+    if 'maps' in path:
+        image_a = image[:, w/2:, :].astype(np.float32) / 255.0
+        image_b = image[:, :w/2, :].astype(np.float32) / 255.0
+    else:
+        image_a = image[:, :w/2, :].astype(np.float32) / 255.0
+        image_b = image[:, w/2:, :].astype(np.float32) / 255.0
 
     # range of pixel values = [-1.0, 1.0]
     image_a = image_a * 2.0 - 1.0