update before release

Author: Ruilong Li

demo.py

@@ -3,21 +3,26 @@
 import torch
 import numpy as np
 import cv2
+import glob
 import streamer_pytorch as streamer
 
+
 parser = argparse.ArgumentParser(description='.')
 parser.add_argument(
-    '--camera', action="store_true")
+    '--camera', action="store_true", help="whether to use webcam.")
+parser.add_argument(
+    '--images', default="", nargs="*", help="paths of image.")
 parser.add_argument(
-    '--images', default="", nargs="*")
+    '--image_folder', default="", help="path of image folder.")
 parser.add_argument(
-    '--videos', default="", nargs="*")
+    '--videos', default="", nargs="*", help="paths of video.")
 parser.add_argument(
-    '--loop', action="store_true")
+    '--loop', action="store_true", help="whether to repeat images/video.")
 parser.add_argument(
-    '--vis', action="store_true")
+    '--vis', action="store_true", help="whether to visualize.")
 args = parser.parse_args()
 
+
 def visulization(data):
     window = data[0].numpy()
     window = window.transpose(1, 2, 0)
@@ -29,14 +34,21 @@ def visulization(data):
     cv2.imshow('window', window)
     cv2.waitKey(1)
 
+
 if args.camera:
-    data_stream = streamer.CaptureStreamer()
+    data_stream = streamer.CaptureStreamer(pad=False)
 elif len(args.videos) > 0:
     data_stream = streamer.VideoListStreamer(
-        args.videos * (100 if args.loop else 1))
+        args.videos * (10 if args.loop else 1))
 elif len(args.images) > 0:
     data_stream = streamer.ImageListStreamer(
-        args.images * (100 if args.loop else 1))
+        args.images * (10000 if args.loop else 1))
+elif args.image_folder is not None:
+    images = sorted(glob.glob(args.image_folder+'/*.jpg'))
+    images += sorted(glob.glob(args.image_folder+'/*.png'))
+    data_stream = streamer.ImageListStreamer(
+        images * (10 if args.loop else 1))
+
 
 loader = torch.utils.data.DataLoader(
     data_stream, 
@@ -45,6 +57,7 @@ def visulization(data):
     pin_memory=False,
 )
 
+
 try:
     for data in tqdm.tqdm(loader):
         if args.vis:

setup.py

@@ -4,9 +4,9 @@
 
 setuptools.setup(
     name='streamer_pytorch',
-    url='https://github.com/liruilong940607/streamer_pytorch',
+    url='https://github.com/Project-Splinter/streamer_pytorch',
     description='Pytorch based data streamer. (Capture, Video & Image).',    
-    version='0.0.1',
+    version='0.0.2',
     author='Ruilong Li',
     author_email='[email protected]',    
     license='MIT License',

streamer_pytorch/streamer.py

@@ -43,6 +43,8 @@ def get_affine_matrix(center, translate, scale):
 
 
 class BaseStreamer():
+    """This streamer will return images at 512x512 size.
+    """
     def __init__(self, 
                  width=512, height=512, pad=True, 
                  mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
@@ -76,6 +78,8 @@ def __len__(self):
 
 
 class CaptureStreamer(BaseStreamer):
+    """This streamer takes webcam as input.
+    """
     def __init__(self, id=0, width=512, height=512, pad=True, **kwargs):
         super().__init__(width, height, pad, **kwargs)
         self.capture = cv2.VideoCapture(id)
@@ -94,6 +98,8 @@ def __del__(self):
 
 
 class VideoListStreamer(BaseStreamer):
+    """This streamer takes a list of video files as input.
+    """
     def __init__(self, files, width=512, height=512, pad=True, **kwargs):
         super().__init__(width, height, pad, **kwargs)
         self.files = files
@@ -115,6 +121,8 @@ def __del__(self):
 
 
 class ImageListStreamer(BaseStreamer):
+    """This streamer takes a list of image files as input.
+    """
     def __init__(self, files, width=512, height=512, pad=True, **kwargs):
         super().__init__(width, height, pad, **kwargs)
         self.files = files
@@ -129,53 +137,4 @@ def __len__(self):
         return len(self.files)
 
 
-if __name__ == "__main__":
-    import tqdm
-    import argparse
-
-    parser = argparse.ArgumentParser(description='.')
-    parser.add_argument(
-        '--camera', action="store_true")
-    parser.add_argument(
-        '--images', default="", nargs="*")
-    parser.add_argument(
-        '--videos', default="", nargs="*")
-    parser.add_argument(
-        '--loop', action="store_true")
-    args = parser.parse_args()
-
-    def visulization(data):
-        window = data[0].numpy()
-        window = window.transpose(1, 2, 0)
-        window = (window * 0.5 + 0.5) * 255.0
-        window = np.uint8(window)
-        window = cv2.cvtColor(window, cv2.COLOR_BGR2RGB) 
-        window = cv2.resize(window, (0, 0), fx=2, fy=2)
-
-        cv2.imshow('window', window)
-        cv2.waitKey(1)
-
-    if args.camera:
-        data_stream = CaptureStreamer()
-    elif len(args.videos) > 0:
-        data_stream = VideoListStreamer(args.videos * (100 if args.loop else 1))
-    elif len(args.images) > 0:
-        data_stream = ImageListStreamer(args.images * (100 if args.loop else 1))
-
-    loader = torch.utils.data.DataLoader(
-        data_stream, 
-        batch_size=1, 
-        num_workers=1, 
-        pin_memory=False,
-    )
-
-    try:
-        for data in tqdm.tqdm(loader):
-            visulization(data)
-            pass
-    except Exception as e:
-        print (e)
-        del data_stream
-
-