improve model uploader ui

app.py

@@ -8,7 +8,6 @@ from PIL import Image
 spec = importlib.util.find_spec('mmdet')
 if spec and spec.origin:
     src = open(spec.origin, encoding='utf-8').read()
-    # strip out the mmcv_minimum_version…assert… block up to __all__
     patched = re.sub(r'(?ms)^[ \t]*mmcv_minimum_version.*?^__all__', '__all__', src)
     m = importlib.util.module_from_spec(spec)
     m.__loader__ = spec.loader
@@ -28,16 +27,18 @@ def load_inferencer(checkpoint_path=None, device=None):
 
 # ——— Gradio prediction function ———
 @spaces.GPU()
-def predict(image: Image.Image, checkpoint: str = None):
+def predict(image: Image.Image, checkpoint):
     # save upload to temp file
     inp_path = "/tmp/upload.jpg"
     image.save(inp_path)
 
+    # determine checkpoint path if user uploaded
+    ckpt_path = checkpoint.name if checkpoint else None
+
     vis_dir = "/tmp/vis"
     os.makedirs(vis_dir, exist_ok=True)
 
-    inferencer = load_inferencer(checkpoint_path=checkpoint, device=None)
-    # run inference & visualization
+    inferencer = load_inferencer(checkpoint_path=ckpt_path, device=None)
     for result in inferencer(
         inputs=inp_path,
         bbox_thr=0.1,
@@ -48,7 +49,6 @@ def predict(image: Image.Image, checkpoint: str = None):
     ):
         pass
 
-    # return the first visualization
     out_files = sorted(os.listdir(vis_dir))
     if out_files:
         return Image.open(os.path.join(vis_dir, out_files[0]))
@@ -59,11 +59,11 @@ demo = gr.Interface(
     fn=predict,
     inputs=[
         gr.Image(type="pil", label="Upload Image"),
-        gr.Textbox(label="RTMO PyTorch Checkpoint Path (optional)")
+        gr.File(file_types=['.pth'], label="Upload RTMO .pth Checkpoint (optional)")
     ],
     outputs=gr.Image(type="pil", label="Annotated Image"),
     title="RTMO Pose Demo",
-    description="Upload an image, optionally supply a RTMO .pth checkpoint, and see 2D pose annotation.",
+    description="Upload an image and (optionally) a RTMO .pth checkpoint to get 2D pose annotation.",
 )
 
 def main():