Add filter file function and test script

main.py

@@ -1,3 +1,4 @@
+from pathlib import Path
 from typing import List
 
 import cv2
@@ -20,6 +21,80 @@ BILATERAL_MODEL.eval()
 
 OUTPUT_GALLERY = gr.Gallery(
     label='Highlighted Area').style(grid=[2], height='auto')
+SUPPORTED_IMG_EXT = ['.png', '.jpg', '.jpeg']
+
+
+def filter_files(files: List) -> List:
+    """Filter uploaded files.
+
+    The model requires a pair of CC-MLO view of the breast scan.
+    This function will filter and ensure the inputs are as expected.
+    FIlter:
+    - Not enough number of files
+    - Unsupported extensions
+    - Missing required pair or part
+
+    Parameters
+    ----------
+    files : List[tempfile._TemporaryFileWrapper]
+        List of path to downloaded files
+
+    Returns
+    -------
+    List[tempfile._TemporaryFileWrapper]
+        List of path to downloaded files
+
+    Raises
+    ------
+    gr.Error
+        If the files is not equal to 2,
+    gr.Error
+        If the extension is unsupported
+    gr.Error
+        If specific view or side of mammography is missing.
+    """
+    if len(files) != 2:
+        raise gr.Error(
+            f'Need exactly 2 images. Currently have {len(files)} images!')
+
+    file_paths = [Path(file.name) for file in files]
+
+    if not all([path.suffix in SUPPORTED_IMG_EXT for path in file_paths]):
+        raise gr.Error(f'There is a file with unsupported type. \
+            Make sure all files are in {SUPPORTED_IMG_EXT}!')
+
+    # Table to store view(row), side(column)
+    table = np.zeros((2, 2), dtype=bool)
+    bin_left = 0
+    bin_right = 0
+    for file in file_paths:
+
+        splits = file.name.split('_')
+
+        # Check if view is present
+        if any(['cc' in part.lower() for part in splits]):
+            table[0, :] = [True, True]
+        if any(['mlo' in part.lower() for part in splits]):
+            table[1, :] = [True, True]
+
+        # Check if side is present
+        if any(['left' in part.lower() for part in splits]):
+            table[:, 0] &= True
+            bin_left += 1
+        elif any(['right' in part.lower() for part in splits]):
+            table[:, 1] &= True
+            bin_right += 1
+
+    # Reset side that has not enough files
+    if bin_left < 2:
+        table[:, 0] &= False
+    if bin_right < 2:
+        table[:, 1] &= False
+
+    if not any([all(table[:, 0]), all(table[:, 1])]):
+        raise gr.Error('Missing bilateral-view pair for Left or Right side.')
+
+    return file_paths
 
 
 def predict_bilateral(file: str) -> List:
@@ -87,7 +162,7 @@ def run():
         outputs=[OUTPUT_GALLERY, gr.Label(label='Cancer Type')]
     )
 
-    demo.launch(server_name='0.0.0.0', server_port=7860)
+    demo.launch(server_name='0.0.0.0', server_port=7860)  # nosec B104
     demo.close()
 
 

tests/test_file_filters.py

@@ -0,0 +1,74 @@
+from pathlib import Path
+from typing import List
+
+import gradio as gr
+import numpy as np
+import pytest
+
+SUPPORTED_IMG_EXT = ['.png', '.jpg', '.jpeg']
+
+
+def filter_files(files: List[str]):
+    if len(files) != 2:
+        raise gr.Error(
+            f'Need exactly 2 images. Currently have {len(files)} images!')
+
+    file_paths = [Path(file) for file in files]
+
+    if not all([path.suffix in SUPPORTED_IMG_EXT for path in file_paths]):
+        raise gr.Error(f'There is a file with unsupported type. \
+            Make sure all files are in {SUPPORTED_IMG_EXT}!')
+
+    # Table to store view(row), side(column)
+    table = np.zeros((2, 2), dtype=bool)
+    bin_left = 0
+    bin_right = 0
+    for file in file_paths:
+
+        splits = file.name.split('_')
+
+        # Check if view is present
+        if any(['cc' in part.lower() for part in splits]):
+            table[0, :] = [True, True]
+        if any(['mlo' in part.lower() for part in splits]):
+            table[1, :] = [True, True]
+
+        # Check if side is present
+        if any(['left' in part.lower() for part in splits]):
+            table[:, 0] &= True
+            bin_left += 1
+        elif any(['right' in part.lower() for part in splits]):
+            table[:, 1] &= True
+            bin_right += 1
+
+    # Reset side that has not enough files
+    if bin_left < 2:
+        table[:, 0] &= False
+    if bin_right < 2:
+        table[:, 1] &= False
+
+    if not any([all(table[:, 0]), all(table[:, 1])]):
+        raise gr.Error('Missing bilateral-view pair for Left or Right side.')
+
+    return file_paths
+
+
+@pytest.mark.parametrize('file_paths,expectation',
+                         [
+                             (['left_cc.png'], False),
+                             (['left_cc.png', 'right_cc.png', 'left_mlo.png'], False),
+                             (['left_cc.png', 'left_mlo.tiff'], False),
+                             (['left_cc.png', 'left_cc.png'], False),
+                             (['left_mlo.png', 'right_cc.png'], False),
+                             (['no_mlo.png', 'left_cc.png'], False),
+                             (['left_cc.png', 'left_mlo.png'], True),
+                             (['right_cc.jpg', 'right_mlo.png'], True),
+                         ]
+                         )
+def test_filter_files(file_paths, expectation):
+    if expectation:
+        filtered_paths = filter_files(file_paths)
+        assert len(filtered_paths) == 2
+    else:
+        with pytest.raises(Exception) as _:
+            _ = filter_files(file_paths)