add: convert_data_webdataset to transform an existing wds file

Author: yztxwd

seqchromloader/__init__.py

@@ -1,2 +1,2 @@
 from .loader import SeqChromDatasetByDataFrame, SeqChromDatasetByBed, SeqChromDatasetByWds, SeqChromDataModule
-from .writer import dump_data_webdataset
+from .writer import dump_data_webdataset, convert_data_webdataset

seqchromloader/loader.py

@@ -57,12 +57,13 @@ class _SeqChromDatasetByWds(IterableDataset):
     :param transforms: A dictionary of functions to transform the output data, accepted keys are **["seq", "chrom", "target", "label"]**
     :type transforms: dict of functions
     """
-    def __init__(self, wds, transforms:dict=None, rank=0, world_size=1):
+    def __init__(self, wds, transforms:dict=None, rank=0, world_size=1, keep_key=False):
         self.wds = wds
         self.transforms = transforms
 
         self.rank = rank
         self.world_size = world_size
+        self.keep_key = keep_key
 
     def initialize(self):
         # this function will be called by worker_init_function in DataLoader
@@ -85,7 +86,10 @@ def __iter__(self):
         if self.transforms is not None: 
             pipeline.append(wds.map_dict(**self.transforms))
 
-        pipeline.append(wds.to_tuple("seq", "chrom", "target", "label"))
+        if self.keep_key:
+            pipeline.append(wds.to_tuple("__key__", "seq", "chrom", "target", "label"))
+        else:
+            pipeline.append(wds.to_tuple("seq", "chrom", "target", "label"))
 
         ds = wds.DataPipeline(*pipeline)
 

seqchromloader/writer.py

@@ -17,14 +17,43 @@
 import pyBigWig
 import webdataset as wds
 
-from seqchromloader import utils
+from . import utils
+from .loader import _SeqChromDatasetByWds
 
+def convert_data_webdataset(wds_in, wds_out, transforms=None, compress=False):
+    """
+    Transform the provided webdataset
+    
+    :param wds_in: input webdataset file
+    :type wds_in: string
+    :param wds_out: output webdataset file
+    :type wds_out: string
+    :param transforms: A dictionary of functions to transform the output data, accepted keys are *["seq", "chrom", "target", "label"]*
+    :type transforms: dict of functions
+    :param compress: whether to compress the output file
+    :type compress: boolean
+    """
+    
+    ds = _SeqChromDatasetByWds(wds_in, transforms=transforms, keep_key=True)
+    sink = wds.TarWriter(wds_out, compress=compress)
+    for (key, seq, chrom, target, label) in ds:
+        feature_dict = defaultdict()
+        feature_dict["__key__"] = key
+        
+        feature_dict["seq.npy"] = seq
+        feature_dict["chrom.npy"] = chrom
+        feature_dict["target.npy"] = target
+        feature_dict["label.npy"] = label
+        sink.write(feature_dict)
+    sink.close()
+    
 def dump_data_webdataset(coords, genome_fasta, bigwig_filelist,
                         target_bam=None, 
                         outdir="dataset/", outprefix="seqchrom", 
                         compress=True, 
                         numProcessors=1,
-                        transforms=None):
+                        transforms=None,
+                        DALI=False):
     """
     Given coordinates dataframe, extract the sequence and chromatin signal, save in webdataset format