training: split_exr: added support for subimages

RenderKit · Dec 13, 2023 · 4425d6d · 4425d6d
1 parent cf98936
commit 4425d6d
Showing 1 changed file with 72 additions and 63 deletions.
diff --git a/training/split_exr.py b/training/split_exr.py
@@ -20,73 +20,82 @@ def main():
     ext = '.exr'
   if ext.lower() != '.exr':
     error('image must be EXR')
-  image = oiio.ImageBuf(cfg.input)
-  if image.has_error:
-    error('could not load image')
 
-  # Get the channels and group them by layer
-  channels = image.spec().channelnames
-  print('Channels:', channels)
-  layer_channels = defaultdict(set)
-  for channel in channels:
-    if len(channel.split('.')) >= 3:
-      layer, ch = channel.split('.', 1)
-      layer_channels[layer].add(ch)
-    else:
-      layer_channels[None].add(channel)
+  # Iterate over subimages
+  subimage = 0
+  while True:
+    image = oiio.ImageBuf(cfg.input, subimage, 0)
+    if image.has_error:
+      error('could not load image')
+    subimage += 1
 
-  # Set default layer
-  if not cfg.layer and len(layer_channels) == 1:
-    cfg.layer = list(layer_channels.keys())[0]
+    # Get the channels and group them by layer
+    channels = image.spec().channelnames
+    if not channels:
+      break
+    print('Channels:', channels)
+    layer_channels = defaultdict(set)
+    for channel in channels:
+      if len(channel.split('.')) >= 3:
+        layer, ch = channel.split('.', 1)
+        layer_channels[layer].add(ch)
+      else:
+        layer_channels[None].add(channel)
 
-  # Extract features
-  FEATURES = {
-    'hdr' : [
-              ('R', 'G', 'B'),
-              ('Noisy Image.R', 'Noisy Image.G', 'Noisy Image.B'),
-              ('Beauty.R', 'Beauty.G', 'Beauty.B'),
-              ('Combined.R', 'Combined.G', 'Combined.B'),
-              ('Composite.Combined.R', 'Composite.Combined.G', 'Composite.Combined.B')
-            ],
-    'a' : [('A',)],
-    'alb' : [
-              ('albedo.R', 'albedo.G', 'albedo.B'),
-              ('Denoising Albedo.R', 'Denoising Albedo.G', 'Denoising Albedo.B'),
-              ('ViewLayer.Denoising Albedo.R', 'ViewLayer.Denoising Albedo.G', 'ViewLayer.Denoising Albedo.B'),
-              ('VisibleDiffuse.R', 'VisibleDiffuse.G', 'VisibleDiffuse.B'),
-              ('diffuse.R', 'diffuse.G', 'diffuse.B'),
-              ('DiffCol.R', 'DiffCol.G', 'DiffCol.B'),
-            ],
-    'nrm' : [
-              ('normal.R', 'normal.G', 'normal.B'),
-              ('normal.X', 'normal.Y', 'normal.Z'),
-              ('N.R', 'N.G', 'N.B'),
-              ('Denoising Normal.X', 'Denoising Normal.Y', 'Denoising Normal.Z'),
-              ('ViewLayer.Denoising Normal.X', 'ViewLayer.Denoising Normal.Y', 'ViewLayer.Denoising Normal.Z'),
-              ('Normals.R', 'Normals.G', 'Normals.B'),
-              ('VisibleNormals.R', 'VisibleNormals.G', 'VisibleNormals.B'),
-              ('OptixNormals.R', 'OptixNormals.G', 'OptixNormals.B'),
-            ],
-    'z' : [
-            ('Denoising Depth.Z',),
-            ('ViewLayer.Denoising Depth.Z',)
-          ]
-  }
+    # Set default layer
+    if not cfg.layer and len(layer_channels) == 1:
+      cfg.layer = list(layer_channels.keys())[0]
 
-  for feature, feature_channel_lists in FEATURES.items():
-    for feature_channels in feature_channel_lists:
-      # Check whether the feature is present in the selected layer of the image
-      if cfg.layer:
-        feature_channels = tuple([cfg.layer + '.' + f for f in feature_channels])
-      if set(feature_channels).issubset(channels):
-        # Save the feature image
-        feature_filename = name + '.' + feature + ext
-        print(feature_filename)
-        new_channels = ('R', 'G', 'B') if len(feature_channels) == 3 else ('Y',)
-        feature_image = oiio.ImageBufAlgo.channels(image, feature_channels, new_channels)
-        feature_image.spec().attribute('compression', 'piz')
-        feature_image.write(feature_filename)
-        break
+    # Extract features
+    FEATURES = {
+      'hdr' : [
+                ('R', 'G', 'B'),
+                ('Noisy Image.R', 'Noisy Image.G', 'Noisy Image.B'),
+                ('Beauty.R', 'Beauty.G', 'Beauty.B'),
+                ('Combined.R', 'Combined.G', 'Combined.B'),
+                ('Composite.Combined.R', 'Composite.Combined.G', 'Composite.Combined.B')
+              ],
+      'a' : [('A',)],
+      'alb' : [
+                ('albedo.R', 'albedo.G', 'albedo.B'),
+                ('Denoising Albedo.R', 'Denoising Albedo.G', 'Denoising Albedo.B'),
+                ('ViewLayer.Denoising Albedo.R', 'ViewLayer.Denoising Albedo.G', 'ViewLayer.Denoising Albedo.B'),
+                ('VisibleDiffuse.R', 'VisibleDiffuse.G', 'VisibleDiffuse.B'),
+                ('diffuse.R', 'diffuse.G', 'diffuse.B'),
+                ('DiffCol.R', 'DiffCol.G', 'DiffCol.B'),
+                ('albedo.red', 'albedo.green', 'albedo.blue'),
+              ],
+      'nrm' : [
+                ('normal.R', 'normal.G', 'normal.B'),
+                ('normal.X', 'normal.Y', 'normal.Z'),
+                ('N.R', 'N.G', 'N.B'),
+                ('Denoising Normal.X', 'Denoising Normal.Y', 'Denoising Normal.Z'),
+                ('ViewLayer.Denoising Normal.X', 'ViewLayer.Denoising Normal.Y', 'ViewLayer.Denoising Normal.Z'),
+                ('Normals.R', 'Normals.G', 'Normals.B'),
+                ('VisibleNormals.R', 'VisibleNormals.G', 'VisibleNormals.B'),
+                ('OptixNormals.R', 'OptixNormals.G', 'OptixNormals.B'),
+                ('normal.red', 'normal.green', 'normal.blue'),
+              ],
+      'z' : [
+              ('Denoising Depth.Z',),
+              ('ViewLayer.Denoising Depth.Z',)
+            ]
+    }
+
+    for feature, feature_channel_lists in FEATURES.items():
+      for feature_channels in feature_channel_lists:
+        # Check whether the feature is present in the selected layer of the image
+        if cfg.layer:
+          feature_channels = tuple([cfg.layer + '.' + f for f in feature_channels])
+        if set(feature_channels).issubset(channels):
+          # Save the feature image
+          feature_filename = name + '.' + feature + ext
+          print(feature_filename)
+          new_channels = ('R', 'G', 'B') if len(feature_channels) == 3 else ('Y',)
+          feature_image = oiio.ImageBufAlgo.channels(image, feature_channels, new_channels)
+          feature_image.spec().attribute('compression', 'piz')
+          feature_image.write(feature_filename)
+          break
 
 if __name__ == '__main__':
   main()