Added new parameter 'compute_key', that allows users to override the …

…default function for computing a sample key. Added information on this parameter to the README.
rom1504 · edwardguil · Jan 14, 2024 · Jan 14, 2024 · Jan 14, 2024 · Jan 14, 2024
commit 2daf299918fc05e17abf6ac4ce2f29f1e7f522ea
diff --git a/README.md b/README.md
@@ -83,6 83,8 @@ with each number being the position in the list. The subfolders avoids having to
 
 If **captions** are provided, they will be saved as 0.txt, 1.txt, ...
 
 If **compute_key** is provided, the filename keys will be computed from this function instead of the default 000000000.jpg to 'unique_key'.jpg
 
 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:
@@ -177,11 179,43 @@ This module exposes a single function `download` which takes the same arguments
 * **max_shard_retry** Number of time to retry failed shards at the end (default *1*)
 * **user_agent_token** Additional identifying token that will be added to the User-Agent header sent with HTTP requests to download images; for example: "img2downloader". (default *None*)
 * **disallowed_header_directives** List of X-Robots-Tags header directives that, if present in HTTP response when downloading an image, will cause the image to be excluded from the output dataset. To ignore x-robots-tags, pass '[]'. (default '["noai", "noimageai", "noindex", "noimageindex"]')
 * **compute_key** A function reference to override the default function used to compute the key for a given sample set. If set to None, img2dataset create keys as a combination of its shard number and count within the shard e.g. 100001.jpg, 100001.txt. (default *None*)
 
 ## 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 img2dataset will complete the download.
 
 ## Compute key paramater function
 
 To override the default method in calculating a samples key, pass a function reference to the compute_key parameter. This function will be passed 6 parameters, and it should return a single string, *unique across the entire dataset* being downloaded. The parameters it will be passed are:
 - key: the index of the sample in the shard
 - shard_id: the shard id this sample belongs too
 - oom_sample_per_shard: the number of samples per shard
 - oom_shard_count: the total number of shards
 - additional_columns: a dictionary containing any additional columns specified in initial parameters for this sample
 
 As an example you can reconstruct the default method for computing a key:
 ```python
 def compute_key(key, shard_id, oom_sample_per_shard, oom_shard_count, additional_columns):
  true_key = (10**oom_sample_per_shard) * shard_id key
  key_format = oom_sample_per_shard oom_shard_count
  str_key = "{true_key:0{key_format}d}".format(
  key_format=key_format, true_key=true_key
  )
  return str_key
 ```
 Alternatively, if your dataset had some additional which you specified, one of which was a uid across the dataset, you could simply do the following:
 ```python
 def compute_key(key, shard_id, oom_sample_per_shard, oom_shard_count, additional_columns):
  return str(additional_columns['uid'])
 ```
 Which would change the output to be:
 * output_folder
  * 00000
  * *your_uid_0*.jpg
  * *your_uid_0*.txt
  * *your_uid_1*.jpg
  ...
 ## Output format choice
 
 Img2dataset support several formats. There are trade off for which to choose:

diff --git a/img2dataset/downloader.py b/img2dataset/downloader.py
@@ -97,6 97,7 @@ def __init__(
  user_agent_token,
  disallowed_header_directives,
  blurring_bbox_col=None,
  compute_key=None,
  ) -> None:
  self.sample_writer_class = sample_writer_class
  self.resizer = resizer
@@ -119,6 120,7 @@ def __init__(
  else {directive.strip().lower() for directive in disallowed_header_directives}
  )
  self.blurring_bbox_col = blurring_bbox_col
  self.compute_key = compute_key
 
  def __call__(
  self,
@@ -213,14 215,20 @@ def data_generator():
  ):
  try:
  _, sample_data = shard_to_dl[key]
- str_key = compute_key(key, shard_id, oom_sample_per_shard, self.oom_shard_count)
- meta = {
- # Skip columns containing a the verification hash and only save the compute hash
  additional_columns = { 
  # Skip columns containing a the verification hash and only save the compute hash 
  **{
  self.column_list[i]: sample_data[i]
  for i in range(len(self.column_list))
  if (hash_indice is None or i != hash_indice)
- },
  }
  }
  if self.compute_key is None:
  str_key = compute_key(key, shard_id, oom_sample_per_shard, self.oom_shard_count, additional_columns)
  else:
  str_key = self.compute_key(key, shard_id, oom_sample_per_shard, self.oom_shard_count, additional_columns)
  meta = {
  **additional_columns, 
  "key": str_key,
  "status": None,
  "error_message": error_message,

diff --git a/img2dataset/main.py b/img2dataset/main.py
@@ -108,6 108,7 @@ def download(
  max_shard_retry: int = 1,
  user_agent_token: Optional[str] = None,
  disallowed_header_directives: Optional[List[str]] = None,
  compute_key: Optional[callable] = None,
 ):
  """Download is the main entry point of img2dataset, it uses multiple processes and download multiple files"""
  if disallowed_header_directives is None:
@@ -247,6 248,7 @@ def signal_handler(signal_arg, frame): # pylint: disable=unused-argument
  user_agent_token=user_agent_token,
  disallowed_header_directives=disallowed_header_directives,
  blurring_bbox_col=bbox_col,
  compute_key=compute_key,
  )
 
  print("Starting the downloading of this file")