refactor (#29)

* wip refactor

* runs now

* fix lint

* test only 3.8

* fake test

* readme

* small fixes

.github/workflows/ci.yml

@@ -31,7 +31,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.6, 3.7, 3.8]
+        python-version: [3.8]
 
     steps:
     - uses: actions/checkout@v2

.gitignore

@@ -7,4 +7,5 @@ __pycache__
 .env*
 wandb
 *.pex
-.pexing
+.pexing
+**/dataset/*

README.md

@@ -11,20 +11,149 @@ Checkout the [design doc](https://docs.google.com/document/d/1_TD2KQLkEegszq4Eip
 
 pip install video2dataset
 
+## Examples
+
+
+## Usage
+
+First get some video url list. For example:
+```
+echo 'https://www.youtube.com/watch?v=0WfKzVqdQqo' >> myvidlist.txt
+```
+
+Then, run the tool:
+
+```
+video2dataset --url_list=myvidlist.txt --output_folder=output_folder
+```
+
+The tool will then automatically download the urls and store them with that format:
+* output_folder
+    * 00000
+        * 000000000.mp4
+        * 000000001.mp4
+        * 000000002.mp4
+
+or as this format if choosing webdataset:
+* output_folder
+    * 00000.tar containing:
+        * 000000000.mp4
+        * 000000001.mp4
+        * 000000002.mp4
+
+with each number being the position in the list. The subfolders avoids having too many files in a single folder.
+
+If **captions** are provided, they will be saved as 0.txt, 1.txt, ...
+
+This can then easily be fed into machine learning training or any other use case.
+
+Also .json files named 0.json, 1.json,... are saved with these keys:
+* url
+* caption
+* key of the form 000010005 : the first 5 digits are the shard id, the last 4 are the index in the shard
+* status : whether the download succeeded
+* error_message
+
+Also a .parquet file will be saved with the same name as the subfolder/tar files containing these same metadata.
+It can be used to analyze the results efficiently.
+
+.json files will also be saved with the same name suffixed by _stats, they contain stats collected during downloading (download time, number of success, ...)
+
 ## Python examples
 
-Checkout these examples to call this as a lib:
-* [example.py](examples/example.py)
 
 ## API
 
-This module exposes a single function `hello_world` which takes the same arguments as the command line tool:
+This module exposes a single function `download` which takes the same arguments as the command line tool:
+
+* **url_list** A file with the list of url of images to download. It can be a folder of such files. (*required*)
+* **output_folder** The path to the output folder. (default *"images"*)
+* **processes_count** The number of processes used for downloading the pictures. This is important to be high for performance. (default *1*)
+* **encode_format** encode format (default *mp4*)
+* **output_format** decides how to save pictures (default *files*)
+  * **files** saves as a set of subfolder containing pictures
+  * **webdataset** saves as tars containing pictures
+  * **parquet** saves as parquet containing pictures as bytes
+  * **tfrecord** saves as tfrecord containing pictures as bytes
+  * **dummy** does not save. Useful for benchmarks
+* **input_format** decides how to load the urls (default *txt*)
+  * **txt** loads the urls as a text file of url, one per line
+  * **csv** loads the urls and optional caption as a csv
+  * **tsv** loads the urls and optional caption as a tsv
+  * **tsv.gz** loads the urls and optional caption as a compressed (gzip) tsv.gz
+  * **json** loads the urls and optional caption as a json
+  * **parquet** loads the urls and optional caption as a parquet
+* **url_col** the name of the url column for parquet and csv (default *url*)
+* **caption_col** the name of the caption column for parquet and csv (default *None*)
+* **number_sample_per_shard** the number of sample that will be downloaded in one shard (default *10000*)
+* **save_additional_columns** list of additional columns to take from the csv/parquet files and save in metadata files (default *None*)
+* **timeout** maximum time (in seconds) to wait when trying to download an image (default *10*)
+* **enable_wandb** whether to enable wandb logging (default *False*)
+* **wandb_project** name of W&B project used (default *video2dataset*)
+* **oom_shard_count** the order of magnitude of the number of shards, used only to decide what zero padding to use to name the shard files (default *5*)
+* **distributor** choose how to distribute the downloading (default *multiprocessing*)
+  * **multiprocessing** use a multiprocessing pool to spawn processes
+  * **pyspark** use a pyspark session to create workers on a spark cluster (see details below)
+* **subjob_size** the number of shards to download in each subjob supporting it, a subjob can be a pyspark job for example (default *1000*)
+* **incremental_mode** Can be "incremental" or "overwrite". For "incremental", video2dataset will download all the shards that were not downloaded, for "overwrite" video2dataset will delete recursively the output folder then start from zero (default *incremental*)
+
+## Incremental mode
+
+If a first download got interrupted for any reason, you can run again with --incremental "incremental" (this is the default) and using the same output folder , the same number_sample_per_shard and the same input urls, and video2dataset will complete the download.
+
+## Output format choice
+
+video2dataset support several formats. There are trade off for which to choose:
+* files: this is the simplest one, images are simply saved as files. It's good for up to 1M samples on a local file system. Beyond that performance issues appear very fast. Handling more than a million files in standard filesystem does not work well.
+* webdataset: webdataset format saves samples in tar files, thanks to [webdataset](https://webdataset.github.io/webdataset/) library, this makes it possible to load the resulting dataset fast in both pytorch, tensorflow and jax. Choose this for most use cases. It works well for any filesystem
+* parquet: parquet is a columnar format that allows fast filtering. It's particularly easy to read it using pyarrow and pyspark. Choose this if the rest of your data ecosystem is based on pyspark. [petastorm](https://github.com/uber/petastorm) can be used to read the data but it's not as easy to use as webdataset
+* tfrecord: tfrecord is a protobuf based format. It's particularly easy to use from tensorflow and using [tf data](https://www.tensorflow.org/guide/data). Use this if you plan to use the dataset only in the tensorflow ecosystem. The tensorflow writer does not use fsspec and as a consequence supports only a limited amount of filesystem, including local, hdfs, s3 and gcs. It is also less efficient than the webdataset writer when writing to other filesystems than local, losing some 30% performance.
+
+## File system support
+
+Thanks to [fsspec](https://filesystem-spec.readthedocs.io/en/latest/), video2dataset supports reading and writing files in [many file systems](https://github.com/fsspec/filesystem_spec/blob/6233f315548b512ec379323f762b70764efeb92c/fsspec/registry.py#L87).
+To use it, simply use the prefix of your filesystem before the path. For example `hdfs://`, `s3://`, `http://`, or `gcs://`.
+Some of these file systems require installing an additional package (for example s3fs for s3, gcsfs for gcs).
+See fsspec doc for all the details.
+
+If you need specific configuration for your filesystem, you may handle this problem by using the [fsspec configuration system](https://filesystem-spec.readthedocs.io/en/latest/features.html#configuration) that makes it possible to create a file such as `.config/fsspec/s3.json` and have information in it such as:
+```
+{
+  "s3": {
+    "client_kwargs": {
+            "endpoint_url": "https://some_endpoint",
+            "aws_access_key_id": "your_user",
+           "aws_secret_access_key": "your_password"
+    }
+  }
+}
+```
+Which may be necessary if using s3 compatible file systems such as [minio](https://min.io/). That kind of configuration also work for all other fsspec-supported file systems.
+
+## Distribution modes
+
+video2dataset supports several distributors.
+* multiprocessing which spawns a process pool and use these local processes for downloading
+* pyspark which spawns workers in a spark pool to do the downloading
+
+multiprocessing is a good option for downloading on one machine, and as such it is the default.
+Pyspark lets video2dataset use many nodes, which makes it as fast as the number of machines.
+It can be particularly useful if downloading datasets with more than a billion image.
+
+### pyspark configuration
+
+In order to use video2dataset with pyspark, you will need to do this:
+1. `pip install pyspark`
+2. use the `--distributor pyspark` option
+3. tweak the `--subjob_size 1000` option: this is the number of images to download in each subjob. Increasing it will mean a longer time of preparation to put the feather files in the temporary dir, a shorter time will mean sending less shards at a time to the pyspark job.
+
+By default a local spark session will be created.
+You may want to create a custom spark session depending on your specific spark cluster.
 
-* **message** the message to print. (*required*)
 
 ## For development
 
-Either locally, or in [gitpod](https://gitpod.io/#https://github.com/rom1504/video2dataset) (do `export PIP_USER=false` there)
+Either locally, or in [gitpod](https://gitpod.io/#https://github.com/iejMac/video2dataset) (do `export PIP_USER=false` there)
 
 Setup a virtualenv:
 
@@ -47,3 +176,6 @@ make test
 You can use `make black` to reformat the code
 
 `python -m pytest -x -s -v tests -k "dummy"` to run a specific test
+
+## Benchmarks
+

benchmark/run_benchmark.sh

@@ -1,4 +1,12 @@
 #!/bin/bash
 
-mkdir -p dataset
-video2dataset benchmark_vids.parquet --dest="dataset" --output-format="files" --metadata-columns="videoID,title,description,start,end"
+rm -rf dataset
+
+video2dataset --url_list="benchmark_vids.parquet" \
+--input_format="parquet" \
+--output_folder="dataset" \
+--output-format="files" \
+--url_col="videoLoc" \
+--caption_col="title" \
+--save_additional_columns='[videoID,description,start,end]' \
+--enable_wandb=True

requirements-test.txt

@@ -4,3 +4,4 @@ pylint==2.13.4
 pytest-cov==3.0.0
 pytest-xdist==2.5.0
 pytest==7.0.1
+types-requests

requirements.txt

@@ -5,4 +5,7 @@ ffmpeg-python
 yt_dlp
 pyarrow
 fsspec
-webdataset
+webdataset
+wandb
+pandas
+tqdm

setup.py

@@ -22,7 +22,7 @@ def _read_reqs(relpath):
         description="Easily create large video dataset from video urls",
         long_description=long_description,
         long_description_content_type="text/markdown",
-        entry_points={"console_scripts": ["video2dataset=video2dataset.cli:main"]},
+        entry_points={"console_scripts": ["video2dataset=video2dataset.main:main"]},
         author="Maciej Kilian",
         author_email="kilianmaciej6@gmail.com",
         url="https://github.com/iejMac/video2dataset",

tests/test_main.py

@@ -1 +1,2 @@
-import pytest
+def test_hello_world():
+    print("hi")

video2dataset/__init__.py

@@ -1,3 +1,3 @@
 """video2dataset"""
 
-from .video2dataset import video2dataset
+from .main import video2dataset

video2dataset/downloader.py → video2dataset/data_reader.py

@@ -54,10 +54,21 @@ def handle_url(url):
     else:
         print("Warning: Incorrect URL type")
         return None, None, ""
-    
+
     return file.name, file, name
 
 
-class Downloader:
-  def __init__(self):
-    pass
+class VideoDataReader:
+    """Video data reader provide data for a video"""
+
+    def __init__(self) -> None:
+        pass
+
+    def __call__(self, row, timeout, retries):
+        key, url = row
+        file_name, file, _ = handle_url(url)
+        with open(file_name, "rb") as vid_file:
+            vid_bytes = vid_file.read()
+        if file is not None:  # for python files that need to be closed
+            file.close()
+        return key, vid_bytes, None