Upload New File

src/tape_irregularity_classifier/cli.py

+import os
+from os import path
+import sys
+from argparse import ArgumentParser, RawTextHelpFormatter
+from typing import NoReturn
+from rich.console import Console
+
+from mpai_cae_arp.types.irregularity import IrregularityType
+from mpai_cae_arp.files import File, FileType
+
+from tape_irregularity_classifier.lib import (
+    CLASSES,
+    collect_irregularity_images,
+    get_irregularity_file,
+    load_model,
+    MODEL_PATH,
+    predict,
+    prepare_images_for_prediction,
+    PROB,
+    verify_path,
+)
+
+AUDIO_IRREG_FILE = 'AudioAnalyser_IrregularityFileOutput2.json'
+VIDEO_IRREG_FILE = 'VideoAnalyser_IrregularityFileOutput2.json'
+TAPE_IRREG_FILE_1 = 'TapeIrregularityClassifier_IrregularityFileOutput1.json'
+TAPE_IRREG_FILE_2 = 'TapeIrregularityClassifier_IrregularityFileOutput2.json'
+
+
+def get_args() -> tuple[str, str]:
+    """
+    Method to obtain arguments from environment variables or command line.
+    Default environment, ignored if a command line argument is passed.
+
+    Raises
+    ------
+    ValueError
+        If no arguments are passed and environment variables are not set.
+
+    Returns
+    -------
+    tuple[str, str]
+        The working path and the name of the Preservation files, which is key element to retrieve necessary files.
+    """
+    if len(sys.argv) > 1:
+        parser = ArgumentParser(
+            prog='tapeIrregularityClassifier',
+            formatter_class=RawTextHelpFormatter,
+            description='A tool that implements MPAI CAE-ARP Tape Irregularity Classifier Technical Specification.\n'
+                        'By default, the configuration parameters are loaded from ./config/args.yaml file,\n'
+                        'but, alternately, you can pass command line arguments to replace them.'
+        )
+        parser.add_argument('-w','--working-path',help='Specify the Working Path, where all input files are stored',required=True)
+        parser.add_argument('-f','--files-name',help='Specify the name of the Preservation files (without extension)',required=True)
+        args = parser.parse_args()
+        working_path = args.working_path
+        files_name = args.files_name
+    else:
+        working_path = os.getenv('WORKING_DIRECTORY')
+        files_name = os.getenv('FILES_NAME')
+        if any([working_path is None, files_name is None]):
+            raise ValueError("Working path or files name not specified!")
+    return working_path, files_name
+
+
+def exit_with_error(error_message: str, error: int, console: Console) -> NoReturn:
+    console.print(f"[red bold]Error: {error_message}")
+    quit(error)
+
+
+def main():
+    """
+    Execute the Tape Irregularity Classifier workflow:
+    1. Get arguments from environment variables or command line.
+    2. Verify that the environment is conformant to the standard.
+    3. Read the irregularities from VideoAnalyser_IrregularityFileOutput2.json.
+    4. Load the pretrained model.
+    5. Collect the irregularity images to be classified.
+    6. Classify the images.
+    7. Write the results to TapeIrregularityClassifier_IrregularityFileOutput1.json and TapeIrregularityClassifier_IrregularityFileOutput2.json.
+
+    Exceptions are handled and the program exits with the appropriate error code.
+    """
+    console = Console()
+
+    console.print("[bold]Welcome to MPAI CAE-ARP TapeIrregularityClassifier!")
+    console.print(f"You are using Python version: {sys.version}")
+
+    try:
+        working_directory, files_name = get_args()
+        temp_path = verify_path(working_directory, files_name)
+        irregularity_file = get_irregularity_file(path.join(temp_path, VIDEO_IRREG_FILE))
+        audio_irregularity_file = get_irregularity_file(path.join(temp_path, AUDIO_IRREG_FILE))
+        irregularity_file.offset = audio_irregularity_file.offset
+
+        console.print(f"{len(irregularity_file.irregularities)} Irregularities listed in {VIDEO_IRREG_FILE}")
+
+        model = load_model(MODEL_PATH)
+
+        with console.status("Collecting irregularity images :camera:", spinner="dots"):
+            irregularity_images = collect_irregularity_images(irregularity_file.irregularities)
+
+        console.print(f"{len(irregularity_images)} images loaded.")
+        images = prepare_images_for_prediction(irregularity_images)
+        probabilities, label = predict(model, images)
+
+        predicted_labels = []
+        # The following array will correspond to the Irregularities that will be listed in the JSON
+        output_irregularities = []
+
+        for i in range(0, len(irregularity_file.irregularities)):
+            # Retrieve current Irregularity probability
+            irr_prob = probabilities[i][label[i]] * 100
+            predicted_labels.append(CLASSES[label[i]] + ' - ' + str(irr_prob) + ' %')
+            irr = irregularity_file.irregularities[i]
+            # Consider valid Irregularities only those with probability > prob
+            if probabilities[i][label[i]] * 100 > PROB:
+                irr.irregularity_type = IrregularityType(CLASSES[label[i]])
+                output_irregularities.append(irr)
+            else:
+                console.print(f"[cyan]Irregularity number {i+1} with probability {irr_prob:.2f} % discarded")
+                console.print(f"[cyan]    ImageURI: {irr.image_URI}")
+                console.print(f"[cyan]    IrregularityID: {irr.irregularity_ID}")
+                # TODO: Delete corresponding image from temp_path
+                # os.remove(irr['ImageURI'])
+
+        console.print(f"Found {len(output_irregularities)} interesting Irregularities.")
+
+        # TODO: IrregularityFileOutput1 shall present only Irregularities relevant to Tape Audio Restoration
+        console.print(f"Saving {TAPE_IRREG_FILE_1} ...")
+        File(path.join(temp_path, TAPE_IRREG_FILE_1), FileType.JSON)\
+            .write_content(irregularity_file)
+        console.print(f"Saving {TAPE_IRREG_FILE_2} ...")
+        File(path.join(temp_path, TAPE_IRREG_FILE_2), FileType.JSON)\
+            .write_content(irregularity_file)
+
+    except FileNotFoundError as e:
+        exit_with_error(e.strerror, os.EX_NOTFOUND, console)
+    except ValueError as e:
+        exit_with_error(
+            "{}\n{}".format(
+                "Working directory or files name not specified!",
+                "Try -h/--help to know more about Tape Irregularity Classifier usage"),
+            os.EX_USAGE, console)
+    except Exception as e:
+        exit_with_error(e, os.EX_SOFTWARE, console)
+
+    console.print("[green] Success! :tada:")
+    quit(os.EX_OK)
+
+
+if __name__ == '__main__':
+    main()