pybind - Fix external source reader (#199)

* Fix external source reader pybind call

* Add gpu backend support in external_source_reader.py

* Add last batch info for all the video readers

* Move import ciupy to LINE 44

* Fix the image extensions

---------

Co-authored-by: Kiriti Gowda <kiritigowda@gmail.com>
Co-authored-by: Lakshmi Kumar <lakshmi.kumar@amd.com>

rocAL_pybind/amd/rocal/fn.py

@@ -1056,7 +1056,7 @@ def box_iou_matcher(*inputs, anchors, high_threshold=0.5,
     return (box_iou_matcher, [])
 
 
-def external_source(source, device=None, color_format=types.RGB, random_shuffle=False, mode=types.EXTSOURCE_FNAME, max_width=2000, max_height=2000):
+def external_source(source, device=None, color_format=types.RGB, random_shuffle=False, mode=types.EXTSOURCE_FNAME, max_width=2000, max_height=2000, last_batch_policy=types.LAST_BATCH_FILL, last_batch_padded=True):
     """
     External Source Reader - User can pass a iterator or callable source.
     @param source (iterator or callable)                                 The source iterator or callable object.
@@ -1074,7 +1074,7 @@ def external_source(source, device=None, color_format=types.RGB, random_shuffle=
     Pipeline._current_pipeline._external_source_user_given_width = max_width
     Pipeline._current_pipeline._external_source_user_given_height = max_height
     kwargs_pybind = {"rocal_color_format": color_format, "is_output": False, "shuffle": random_shuffle, "loop": False, "decode_size_policy": types.USER_GIVEN_SIZE,
-                     "max_width": max_width, "max_height": max_height, "dec_type": types.DECODER_TJPEG, "external_source_mode": mode}
+                     "max_width": max_width, "max_height": max_height, "dec_type": types.DECODER_TJPEG, "external_source_mode": mode, "last_batch_info": (last_batch_policy, last_batch_padded)}
     external_source_operator = b.externalFileSource(
         Pipeline._current_pipeline._handle, *(kwargs_pybind.values()))
     return (external_source_operator, [])  # Labels is Empty

rocAL_pybind/amd/rocal/readers.py

@@ -189,7 +189,7 @@ def caffe2(path, bbox=False, stick_to_shard=False, pad_last_batch=False):
 def video(sequence_length, file_list_frame_num=False, file_root="", image_type=types.RGB, num_shards=1,
           random_shuffle=False, step=1, stride=1, decoder_mode=types.SOFTWARE_DECODE, enable_frame_num=False,
           enable_timestamps=False, file_list="", stick_to_shard=False, pad_last_batch=False,
-          file_list_include_preceding_frame=False, normalized=False, skip_vfr_check=False):
+          file_list_include_preceding_frame=False, normalized=False, skip_vfr_check=False, last_batch_policy=types.LAST_BATCH_FILL, last_batch_padded=True):
     """!Creates a VideoDecoder node for loading video sequences.
 
         @param sequence_length                      Number of frames in video sequence.
@@ -235,7 +235,8 @@ def video(sequence_length, file_list_frame_num=False, file_root="", image_type=t
         "loop": False,
         "frame_step": step,
         "frame_stride": stride,
-        "file_list_frame_num": file_list_frame_num}  # VideoDecoder
+        "file_list_frame_num": file_list_frame_num,
+        "last_batch_info": (last_batch_policy, last_batch_padded)}  # VideoDecoder
     videos = b.videoDecoder(
         Pipeline._current_pipeline._handle, *(kwargs_pybind_decoder.values()))
     return (videos)
@@ -248,7 +249,8 @@ def video_resize(sequence_length, resize_width, resize_height, file_list_frame_n
                  scaling_mode=types.SCALING_MODE_DEFAULT, interpolation_type=types.LINEAR_INTERPOLATION,
                  resize_longer=0, resize_shorter=0, max_size=[], enable_frame_num=False,
                  enable_timestamps=False, file_list="", stick_to_shard=False, pad_last_batch=False,
-                 file_list_include_preceding_frame=False, normalized=False, skip_vfr_check=False):
+                 file_list_include_preceding_frame=False, normalized=False, skip_vfr_check=False, 
+                 last_batch_policy=types.LAST_BATCH_FILL, last_batch_padded=True):
     """!Creates a VideoDecoderResize node in the pipeline for loading and resizing video sequences.
 
         @param sequence_length                      Number of frames in video sequence.
@@ -294,13 +296,13 @@ def video_resize(sequence_length, resize_width, resize_height, file_list_frame_n
                              "sequence_length": sequence_length, "resize_width": resize_width, "resize_height": resize_height,
                              "shuffle": random_shuffle, "is_output": False, "loop": False, "frame_step": step, "frame_stride": stride,
                              "file_list_frame_num": file_list_frame_num, "scaling_mode": scaling_mode, "max_size": max_size,
-                             "resize_shorter": resize_shorter, "resize_longer": resize_longer, "interpolation_type": interpolation_type}
+                             "resize_shorter": resize_shorter, "resize_longer": resize_longer, "interpolation_type": interpolation_type, "last_batch_info": (last_batch_policy, last_batch_padded)}
     videos = b.videoDecoderResize(
         Pipeline._current_pipeline._handle, *(kwargs_pybind_decoder.values()))
     return (videos, meta_data)
 
 
-def sequence_reader(file_root, sequence_length, image_type=types.RGB, num_shards=1, random_shuffle=False, step=3, stride=1, stick_to_shard=False, pad_last_batch=False):
+def sequence_reader(file_root, sequence_length, image_type=types.RGB, num_shards=1, random_shuffle=False, step=3, stride=1, stick_to_shard=False, pad_last_batch=False, last_batch_policy=types.LAST_BATCH_FILL, last_batch_padded=True):
     """!Creates a SequenceReader node for loading image sequences.
 
         @param file_root            Root directory containing image sequences.
@@ -326,7 +328,8 @@ def sequence_reader(file_root, sequence_length, image_type=types.RGB, num_shards
         "shuffle": random_shuffle,
         "loop": False,
         "frame_step": step,
-        "frame_stride": stride}
+        "frame_stride": stride,
+        "last_batch_info": (last_batch_policy, last_batch_padded)}
     frames = b.sequenceReader(
         Pipeline._current_pipeline._handle, *(kwargs_pybind.values()))
     return (frames)

tests/python_api/external_source_reader.py

@@ -5,16 +5,19 @@
 import amd.rocal.types as types
 import os
 import numpy as np
-import cupy as cp
 import cv2
-from PIL import Image
+import sys
+import glob
 
 
 def main():
-    batch_size = 3
+    if  len(sys.argv) < 3:
+        print ('Please pass cpu/gpu batch_size')
+        exit(0)
+    batch_size = int(sys.argv[2])
+    device = "cpu" if sys.argv[1] == "cpu" else "gpu"
     data_dir = os.environ["ROCAL_DATA_PATH"] + \
-        "rocal_data/coco/coco_10_img/train_10images_2017/"
-    device = "cpu"
+        "rocal_data/coco/coco_10_img/images/"
     try:
         path_mode0 = "output_folder/external_source_reader/mode0/"
         isExist = os.path.exists(path_mode0)
@@ -39,6 +42,7 @@ def main():
 
     def image_dump(img, idx, device="cpu", mode=0):
         if device == "gpu":
+            import cupy as cp
             img = cp.asnumpy(img)
         img = img.transpose([1, 2, 0])  # NCHW
         img = (img).astype('uint8')
@@ -54,9 +58,9 @@ def __init__(self, batch_size):
             self.images_dir = data_dir
             self.batch_size = batch_size
             self.files = []
-            import glob
-            for filename in glob.glob(os.path.join(self.images_dir, '*.jpg')):
-                self.files.append(filename)
+            self.file_patterns = ['*.jpg', '*.jpeg', '*.JPG', '*.JPEG']
+            for pattern in self.file_patterns:
+                self.files.extend(glob.glob(os.path.join(self.images_dir, pattern)))
             shuffle(self.files)
 
         def __iter__(self):
@@ -119,9 +123,9 @@ def __init__(self, batch_size):
             self.batch_size = batch_size
             self.files = []
             import os
-            import glob
-            for filename in glob.glob(os.path.join(self.images_dir, '*.jpg')):
-                self.files.append(filename)
+            self.file_patterns = ['*.jpg', '*.jpeg', '*.JPG', '*.JPEG']
+            for pattern in self.file_patterns:
+                self.files.extend(glob.glob(os.path.join(self.images_dir, pattern)))
 
         def __iter__(self):
             self.i = 0
@@ -190,9 +194,9 @@ def __init__(self, batch_size):
             self.files = []
             self.maxHeight = self.maxWidth = 0
             import os
-            import glob
-            for filename in glob.glob(os.path.join(self.images_dir, '*.jpg')):
-                self.files.append(filename)
+            self.file_patterns = ['*.jpg', '*.jpeg', '*.JPG', '*.JPEG']
+            for pattern in self.file_patterns:
+                self.files.extend(glob.glob(os.path.join(self.images_dir, pattern)))
             shuffle(self.files)
             self.i = 0
             self.n = len(self.files)