Upload decompress_useful.py with huggingface_hub

decompress_useful.py

@@ -0,0 +1,71 @@
+import os
+import glob
+import subprocess
+import argparse
+
+def restore_dataset(input_path="shards", target_path="samples", sensors_to_extract=None):
+    """
+    Restores the dataset from shards.
+    :param sensors_to_extract: Optional list of strings. 
+                               Example: ["LIDAR_TOP", "CAM_FRONT"]
+    """
+    # Ensure the target directory exists
+    if not os.path.exists(target_path):
+        os.makedirs(target_path)
+
+    # Find all unique base names (e.g., "LIDAR_TOP.tar" from "LIDAR_TOP.tar.aa")
+    all_shards = glob.glob(f"{input_path}/*.tar.*")
+    tar_groups = sorted(set(f.split(".tar.")[0] for f in all_shards))
+
+    for base in tar_groups:
+        sensor_name = os.path.basename(base)
+        
+        # 🟢 FILTER LOGIC:
+        # If the user provided a list, skip sensors not in that list
+        if sensors_to_extract:
+            if not any(requested in sensor_name for requested in sensors_to_extract):
+                continue
+
+        md5_file = f"{base}.md5"
+
+        # 1. Verify MD5 integrity if the checksum file exists
+        if os.path.exists(md5_file):
+            print(f"🔍 Verifying integrity for {sensor_name}...")
+            try:
+                # We use shell=True to handle the redirection in the md5sum command
+                subprocess.run(f"md5sum -c {md5_file}", shell=True, check=True)
+                print(f"✨ Checksum verified for {sensor_name}.")
+            except subprocess.CalledProcessError:
+                print(f"⚠️ WARNING: {sensor_name} shards are corrupted! Skipping.")
+                continue
+        else:
+            print(f"ℹ️ No MD5 file found for {sensor_name}, skipping verification.")
+
+        # 2. Reconstruct and extract
+        print(f"🔄 Restoring {sensor_name} to {target_path}...")
+        
+        # Check if the shards are gzipped (ending in .tar.gz.aa) or just tar (.tar.aa)
+        # If 'gz' is in the filename, we add the 'z' flag to tar
+        is_gz = ".gz" in base
+        tar_flags = "-xzf" if is_gz else "-xf"
+        
+        # Command: concatenate all parts and pipe to tar
+        cmd_extract = f"cat {base}.tar.* | tar {tar_flags} - -C {target_path}"
+        
+        try:
+            subprocess.run(cmd_extract, shell=True, check=True)
+            print(f"✅ {sensor_name} successfully restored.")
+        except subprocess.CalledProcessError as e:
+            print(f"❌ Error extracting {sensor_name}: {e}")
+
+if __name__ == "__main__":
+    # Example usage:
+    # To extract everything: restore_dataset()
+    # To extract only LiDAR: restore_dataset(sensors_to_extract=["LIDAR"])
+    
+    # Optional: Integration with command line arguments
+    parser = argparse.ArgumentParser(description="Restore dataset from shards.")
+    parser.add_argument("--sensors", nargs="+", help="List of sensors to extract (e.g. LIDAR NARROW)")
+    args = parser.parse_args()
+
+    restore_dataset(sensors_to_extract=args.sensors)