Lazy load PIL

PiperOrigin-RevId: 200762820

Author: Ryan Sepassi

tensor2tensor/data_generators/bair_robot_pushing.py

@@ -27,7 +27,6 @@
 import os
 import tarfile
 import numpy as np
-from PIL import Image
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -40,6 +39,12 @@
     "http://rail.eecs.berkeley.edu/datasets/bair_robot_pushing_dataset_v0.tar")
 
 
+# Lazy load PIL.Image
+def PIL_Image():  # pylint: disable=invalid-name
+  from PIL import Image  # pylint: disable=g-import-not-at-top
+  return Image
+
+
 @registry.register_problem
 class VideoBairRobotPushing(video_utils.VideoProblem):
   """Berkeley (BAIR) robot pushing dataset."""
@@ -85,7 +90,7 @@ def parse_frames(self, filenames):
           state_name = state_key.format(i)
 
           byte_str = x.features.feature[image_name].bytes_list.value[0]
-          img = Image.frombytes(
+          img = PIL_Image().frombytes(
               "RGB", (self.frame_width, self.frame_height), byte_str)
           arr = np.array(img.getdata())
           frame = arr.reshape(

tensor2tensor/data_generators/google_robot_pushing.py

@@ -27,7 +27,6 @@
 import io
 import os
 import numpy as np
-from PIL import Image
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -42,6 +41,12 @@
 DATA_TEST_NOVEL = (5, "/push_testnovel/push_testnovel.tfrecord-{:05d}-of-00005")
 
 
+# Lazy load PIL.Image
+def PIL_Image():  # pylint: disable=invalid-name
+  from PIL import Image  # pylint: disable=g-import-not-at-top
+  return Image
+
+
 @registry.register_problem
 class VideoGoogleRobotPushing(video_utils.VideoProblem):
   """Google robot pushing dataset."""
@@ -90,10 +95,10 @@ def parse_frames(self, filename):
         state_name = state_key.format(i)
 
         byte_str = x.features.feature[image_name].bytes_list.value[0]
-        img = Image.open(io.BytesIO(byte_str))
+        img = PIL_Image().open(io.BytesIO(byte_str))
         # The original images are much bigger than 64x64
         img = img.resize((self.frame_width, self.frame_height),
-                         resample=Image.BILINEAR)
+                         resample=PIL_Image().BILINEAR)
         arr = np.array(img.getdata())
         frame = arr.reshape(
             self.frame_width, self.frame_height, self.num_channels)