basics work with tarfile as input

config/config.yml

@@ -31,7 +31,7 @@ PATHS:
 VISXP_EXTRACT:
     MODEL_PATH: /model/checkpoint.tar
     MODEL_CONFIG_PATH: /model/model_config.yml
-    TEST_INPUT_PATH: /data
+    TEST_INPUT_PATH: /data/testob/visxp_prep__testob.tar.gz
 INPUT:
     DELETE_ON_COMPLETION: True
 OUTPUT:

feature_extraction.py

@@ -9,7 +9,7 @@
 from data_handling import VisXPData
 from models import VisXPFeatureExtractionOutput
 from provenance import generate_full_provenance_chain
-from io_util import get_source_id, save_features_to_file
+from io_util import get_source_id, save_features_to_file, untar_input_file
 
 logger = logging.getLogger(__name__)
 
@@ -38,12 +38,21 @@ def extract_features(
     source_id = get_source_id(input_path)
     logger.info(f"Extracting features for: {source_id}.")
 
-    # Step 2: Load spectograms + keyframes from file & preprocess
+    # Step 2: check the type of input (tar.gz vs a directory)
+    if input_path.find(".tar.gz") != -1:
+        logger.info("Input is an archive, uncompressing it")
+        untar_input_file(input_path)  # extracts contents in same dir
+        input_path = str(
+            Path(input_path).parent
+        )  # change the input path to the parent dir
+        logger.info(f"Changed input_path to: {input_path}")
+
+    # Step 3: Load spectograms + keyframes from file & preprocess
     dataset = VisXPData(
         Path(input_path), model_config_file=model_config_file, device=device
     )
 
-    # Step 3: Load model from file
+    # Step 4: Load model from file
     model = load_model_from_file(
         checkpoint_file=model_path,
         config_file=model_config_file,
@@ -52,7 +61,7 @@ def extract_features(
     # Switch model mode: in training mode, model layers behave differently!
     model.eval()
 
-    # Step 4: Apply model to data
+    # Step 5: Apply model to data
     logger.info(f"Going to extract features for {dataset.__len__()} items. ")
 
     result_list = []

io_util.py

@@ -1,5 +1,7 @@
 import logging
 import os
+from pathlib import Path
+import tarfile
 from time import time
 import torch
 
@@ -206,3 +208,10 @@ def _fetch_visxp_prep_s3_uri(handler, doc: Document) -> str:
         return possibles[0].payload.get("s3_location", "")
     logger.error("No s3_location found in VISXP_PREP result")
     return ""
+
+
+# untars visxp_prep__<source_id>.tar.gz into the same dir
+def untar_input_file(tar_file_path: str):
+    logger.info(f"Uncompressing {tar_file_path}")
+    with tarfile.open(tar_file_path) as tar:
+        tar.extractall(path=str(Path(tar_file_path).parent), filter="data")  # type: ignore