add spaces for ZeroGPU

app.py

@@ -2,6 +2,7 @@ import cv2
 import gradio as gr
 import json
 import numpy as np
+import spaces
 import torch
 import torch.nn as nn
 
@@ -70,6 +71,8 @@ greulich_and_pyle_ages = {k: np.asarray(v) for k, v in greulich_and_pyle_ages.it
 model_grad_cam = ModelForGradCAM(model_list[0])
 target_layers = [model_grad_cam.model.backbone.stages[-1]]
 
+
+@spaces.GPU
 def predict_bone_age(Radiograph, Sex, Heatmap):
     x0 = rearrange(Radiograph, "h w -> h w 1")
     x = cfg_crop.val_transforms(image=x0)["image"]
@@ -118,18 +121,27 @@ def predict_bone_age(Radiograph, Sex, Heatmap):
     if Heatmap:
         targets = [ClassifierOutputTarget(round(bone_age))]
         with GradCAM(model=model_grad_cam, target_layers=target_layers) as cam:
-            grayscale_cam = cam(input_tensor=x.to(device).float(), targets=targets, eigen_smooth=True)
+            grayscale_cam = cam(
+                input_tensor=x.to(device).float(), targets=targets, eigen_smooth=True
+            )
 
-        heatmap = cv2.applyColorMap((grayscale_cam[0] * 255).astype("uint8"), cv2.COLORMAP_JET)
+        heatmap = cv2.applyColorMap(
+            (grayscale_cam[0] * 255).astype("uint8"), cv2.COLORMAP_JET
+        )
         image = cv2.cvtColor(x[0, 0].cpu().numpy().astype("uint8"), cv2.COLOR_GRAY2RGB)
         image_weight = 0.6
         grad_cam_image = (1 - image_weight) * heatmap[..., ::-1] + image_weight * image
         grad_cam_image = grad_cam_image.astype("uint8")
     else:
         # if no heatmap desired, just show image
-        grad_cam_image = cv2.cvtColor(x[0, 0].cpu().numpy().astype("uint8"), cv2.COLOR_GRAY2RGB)
+        grad_cam_image = cv2.cvtColor(
+            x[0, 0].cpu().numpy().astype("uint8"), cv2.COLOR_GRAY2RGB
+        )
 
-    return f"Predicted bone age: {bone_age_str}\n\nThe closest Greulich & Pyle bone ages are:\n  1) {closest1}\n  2) {closest2}", grad_cam_image
+    return (
+        f"Predicted bone age: {bone_age_str}\n\nThe closest Greulich & Pyle bone ages are:\n  1) {closest1}\n  2) {closest2}",
+        grad_cam_image,
+    )
 
 
 image = gr.Image(image_mode="L")
@@ -170,7 +182,7 @@ with gr.Blocks() as demo:
             ["examples/10043.png", "Female", "No"],
             ["examples/8888.png", "Female", "Yes"],
         ],
-        cache_examples=False
+        cache_examples=False,
     )
 
 if __name__ == "__main__":

requirements.txt

@@ -3,5 +3,6 @@ einops
 grad-cam
 gradio
 scikit-image
+spaces
 timm
 torch