split the functions

app.py

@@ -36,7 +36,7 @@ model.eval()
 
 tf = v2.Compose([
     v2.ToImage(),
-    v2.Resize((64, 64)),
+    v2.Resize((64, 64), antialias=True),
     v2.ToDtype(torch.float32, scale=True),
 ])
 
@@ -74,17 +74,26 @@ for i in range(number_of_examples):
     example_rows_val.append([val_file_path, val_truth_label_text])
 
 @spaces.GPU
-def predict(image,truth_labels=None):
+def predict(image):
     if image is None:
         return None
     
-    if truth_labels:
-        # already preprocessed test/val_samples
-        img_tensor = torch.from_numpy(image).unsqueeze(0).to(device)
-    else:
-        # uplouded image
-        pil_img = Image.fromarray(image.astype('uint8'), 'RGB')
-        img_tensor = tf(pil_img).unsqueeze(0).to(device)
+    # uplouded image
+    img_tensor = tf(image).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        probabilities = torch.nn.functional.softmax(outputs.squeeze(0), dim=0)
+    
+    return {LABELS[i]: float(probabilities[i]) for i in range(len(LABELS))}
+
+@spaces.GPU
+def sample_predict(image,truth_labels=None):
+    if image is None:
+        return None
+    
+    # already preprocessed test/val_samples
+    img_tensor = torch.from_numpy(image.astype(numpy.float32)).unsqueeze(0).to(device)
 
     with torch.no_grad():
         outputs = model(img_tensor)
@@ -101,11 +110,10 @@ custom_css = """
 }
 """
 
-
-with gradio.Blocks(css=custom_css) as demo:
+with gradio.Blocks() as demo:
     gradio.Markdown("# PathMNIST Image Classification")
 
-    with gradio.Tab("Predict"):
+    with gradio.Tab("Predict",css=custom_css):
         gradio.Markdown("## Upload a tissue image for classification")
         with gradio.Row():
             input_img = gradio.Image(height=512, width=512)
@@ -114,7 +122,7 @@ with gradio.Blocks(css=custom_css) as demo:
                 btn = gradio.Button("Predict")
                 btn.click(fn=predict, inputs=input_img, outputs=output_lbl)
 
-    with gradio.Tab("Examples (Validation)"):
+    with gradio.Tab("Examples (Validation)",css=custom_css):
         gradio.Markdown("## Select an example below to test the model against the PathMNIST validation dataset")
         
         with gradio.Row():
@@ -127,11 +135,11 @@ with gradio.Blocks(css=custom_css) as demo:
             examples=example_rows_val, 
             inputs=[input_img_ex, truth_box],
             outputs=output_lbl_ex,
-            fn=predict,
+            fn=sample_predict,
             cache_examples=True,
         )
 
-    with gradio.Tab("Examples (Test)"):
+    with gradio.Tab("Examples (Test)",css=custom_css):
         gradio.Markdown("## Select an example below to test the model against the PathMNIST test dataset")
         
         with gradio.Row():
@@ -144,7 +152,7 @@ with gradio.Blocks(css=custom_css) as demo:
             examples=example_rows_test, 
             inputs=[input_img_ex, truth_box],
             outputs=output_lbl_ex,
-            fn=predict,
+            fn=sample_predict,
             cache_examples=True,
         )