Add main deployment script

main.py

@@ -0,0 +1,95 @@
+from typing import List
+
+import cv2
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+
+from models import phc_models
+from utils import utils
+
+BILATERIAL_WEIGHT = 'weights/phresnet18_cbis2views.pt'
+BILATERAL_MODEL = phc_models.PHCResNet18(
+    channels=2, n=2, num_classes=1, visualize=True)
+BILATERAL_MODEL.add_top_blocks(num_classes=1)
+BILATERAL_MODEL.load_state_dict(torch.load(
+    BILATERIAL_WEIGHT, map_location='cpu'))
+BILATERAL_MODEL = BILATERAL_MODEL.to('cpu')
+BILATERAL_MODEL.eval()
+
+OUTPUT_GALLERY = gr.Gallery(
+    label='Highlighted Area').style(grid=[2], height='auto')
+
+
+def predict_bilateral(file: str) -> List:
+    """Predict Bilateral Mammography.
+
+    Parameters
+    ----------
+    file : TemporaryFileWrapper
+        TemporaryFile object for the uploaded file
+
+    Returns
+    -------
+    List[List, Dict]
+        List of objects that will be used to display the result
+    """
+    displays_imgs = []
+
+    image = np.array(Image.open(file.name))/257
+    image = np.reshape(image, (2, image.shape[0]//2, image.shape[1]))
+
+    im_h, im_w = image[0].shape[:2]
+
+    image_t = torch.from_numpy(image)
+    image_t = image_t.unsqueeze(0)  # Add batch dimension
+
+    out, _, out_refiner = BILATERAL_MODEL(image_t)
+
+    out_refiner = utils.mean_activations(out_refiner).numpy()
+
+    probability = torch.sigmoid(out).detach().cpu().item()
+    label_name = 'Malignant' if probability > 0.5 else 'Normal/Benign'
+    lebels_dict = {label_name: probability}
+
+    refined_view_norm = cv2.normalize(
+        out_refiner, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
+    refined_view = cv2.applyColorMap(refined_view_norm, cv2.COLORMAP_JET)
+    refined_view = cv2.resize(
+        refined_view, (im_w, im_h), interpolation=cv2.INTER_LINEAR)
+
+    image0_colored = cv2.normalize(
+        image[0], None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
+    image0_colored = cv2.cvtColor(image0_colored, cv2.COLOR_GRAY2RGB)
+    image1_colored = cv2.normalize(
+        image[1], None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
+    image1_colored = cv2.cvtColor(image1_colored, cv2.COLOR_GRAY2RGB)
+
+    heatmap0_overlay = cv2.addWeighted(
+        image0_colored, 1.0, refined_view, 0.5, 0)
+    heatmap1_overlay = cv2.addWeighted(
+        image1_colored, 1.0, refined_view, 0.5, 0)
+
+    displays_imgs += [(image0_colored, 'CC'), (image1_colored, 'MLO')]
+
+    displays_imgs.append((heatmap0_overlay, 'CC Interest Area'))
+    displays_imgs.append((heatmap1_overlay, 'MLO Interest Area'))
+
+    return displays_imgs, lebels_dict
+
+
+def run():
+    """Run Gradio App."""
+    demo = gr.Interface(
+        fn=predict_bilateral,
+        inputs=gr.File(file_count='single', file_types=['.png']),
+        outputs=[OUTPUT_GALLERY, gr.Label(label='Cancer Type')]
+    )
+
+    demo.launch(server_name='0.0.0.0', server_port=7860)
+    demo.close()
+
+
+if __name__ == '__main__':
+    run()