Update codes to enable device selection

app.py

@@ -1,11 +1,24 @@
+import json
+import os
 from pathlib import Path
-import torch
-from monai.bundle import ConfigParser
+
 import gradio as gr
 import numpy as np
+import torch
+from monai.bundle import ConfigParser
+
+with open("configs/inference.json") as f:
+    inference_config = json.load(f)
+
+device = torch.device('cpu')
+if torch.cuda.is_available():
+    device = torch.device('cuda:0')
+
+# * NOTE: device must be hardcoded, config file won't affect the device selection
+inference_config["device"] = device
 
 parser = ConfigParser()
-parser.read_config(f="configs/inference.json")
+parser.read_config(f=inference_config)
 parser.read_meta(f="configs/metadata.json")
 
 inference = parser.get_parsed_content("inferer")
@@ -14,9 +27,17 @@ network = parser.get_parsed_content("network_def")
 preprocess = parser.get_parsed_content("preprocessing")
 postprocess = parser.get_parsed_content("postprocessing")
 
+use_fp16 = os.environ.get('USE_FP16', False)
+
 state_dict = torch.load("models/model.pt")
 network.load_state_dict(state_dict, strict=True)
 
+network = network.to(device)
+network.eval()
+
+if use_fp16 and torch.cuda.is_available():
+    network = network.half()
+
 label2color = {0: (0, 0, 0),
              1: (225, 24, 69), # RED
              2: (135, 233, 17), # GREEN
@@ -38,8 +59,11 @@ def visualize_instance_seg_mask(mask):
 def query_image(img):
     data = {"image": img}
     batch = preprocess(data)
+    batch['image'] = batch['image'].to(device)
+
+    if use_fp16 and torch.cuda.is_available():
+        batch['image'] = batch['image'].half()
 
-    network.eval()
     with torch.no_grad():
         pred = inference(batch['image'].unsqueeze(dim=0), network)
 
@@ -65,7 +89,7 @@ with open('Description.md','r') as file:
    markdown_content = file.read()
 
 demo = gr.Interface(
-    query_image, 
+    query_image,
     inputs=[gr.Image(type="filepath")],
     outputs="image",
     title="Medical Image Classification with MONAI - Pathology Nuclei Segmentation Classification",

configs/inference.json

@@ -12,7 +12,7 @@
     "hovernet_mode": "fast",
     "patch_size": 256,
     "out_size": 164,
-    "device": "cpu",
+    "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')",
     "network_def": {
         "_target_": "HoVerNet",
         "mode": "@hovernet_mode",