Upload folder using huggingface_hub

app.py

@@ -1,14 +1,13 @@
-import gradio as gr
-
-def greet(name):
-    return f"Hello, {name}! 👋"
-
-with gr.Blocks() as demo:
-    gr.Markdown("# cspws-space")
-    name = gr.Textbox(label="Your name")
-    out = gr.Textbox(label="Greeting")
-    btn = gr.Button("Say hi")
-    btn.click(fn=greet, inputs=name, outputs=out)
-
-if __name__ == "__main__":
-    demo.launch()
+import os, sys
+from pathlib import Path
+if (Path("scripts") / "gradio_app.py").exists():
+    sys.path.append(str(Path("scripts").absolute()))
+    from gradio_app import main
+    if __name__ == "__main__":
+        os.system("python scripts/gradio_app.py --onnx model.onnx --channels 3 --img-size 256 --port 7860")
+else:
+    import gradio as gr
+    def hello(x): return x
+    demo = gr.Interface(fn=hello, inputs="image", outputs="image", title="csPWS Fallback App")
+    if __name__ == "__main__":
+        demo.launch(server_port=7860)

model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:abe3d43a2d19a4b8eb46d7cbf591c63541b8085c00437cff30cde8797bcbc1d5
+size 39925272

requirements.txt

@@ -1 +1,5 @@
-gradio
+gradio>=4.0.0
+onnxruntime>=1.18.0
+onnx>=1.15.0
+numpy
+pillow

runtime.txt

@@ -0,0 +1 @@
+python-3.12

scripts/gradio_app.py

@@ -0,0 +1,77 @@
+
+import argparse, os, numpy as np, torch
+from PIL import Image
+import torchvision.transforms.functional as F
+import onnxruntime as ort
+
+def pick_providers():
+    avail = ort.get_available_providers()
+    return (["CoreMLExecutionProvider","CPUExecutionProvider"]
+            if "CoreMLExecutionProvider" in avail else ["CPUExecutionProvider"])
+
+def load_torch_model(ckpt_path, channels=3, img_size=256):
+    from cspws_seg.models.unet_attn import UNetAttn
+    device = torch.device("cuda" if torch.cuda.is_available() else ("mps" if (hasattr(torch.backends,'mps') and torch.backends.mps.is_available()) else "cpu"))
+    model = UNetAttn(in_ch=channels, out_ch=1).to(device)
+    state = torch.load(ckpt_path, map_location=device)
+    model.load_state_dict(state.get("model", state)); model.eval()
+    return model, device, img_size
+
+def predict_torch(model, device, img: Image.Image, img_size=256):
+    img = img.convert("RGB")
+    img = F.resize(img, [img_size, img_size], antialias=True)
+    t = F.to_tensor(img).unsqueeze(0).to(device)
+    with torch.inference_mode():
+        logits = model(t)[0,0].detach().cpu().numpy()
+    mask = (1/(1+np.exp(-logits)) > 0.5).astype(np.uint8)*255
+    return Image.fromarray(mask)
+
+def load_onnx_session(onnx_path):
+    prov = pick_providers()
+    sess = ort.InferenceSession(onnx_path, providers=prov)
+    input_name = sess.get_inputs()[0].name
+    return sess, input_name
+
+def predict_onnx(sess, input_name, img: Image.Image, img_size=256):
+    img = img.convert("RGB").resize((img_size,img_size))
+    x = np.transpose(np.array(img, dtype=np.float32)/255.0, (2,0,1))[None,...]
+    logits = sess.run(None, {input_name: x})[0][0,0]
+    mask = (1/(1+np.exp(-logits)) > 0.5).astype(np.uint8)*255
+    return Image.fromarray(mask)
+
+def main():
+    import gradio as gr
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", default="runs/cspws_unet_attn/best.pt")
+    ap.add_argument("--onnx", default=None)
+    ap.add_argument("--img-size", type=int, default=256)
+    ap.add_argument("--port", type=int, default=7860)
+    ap.add_argument("--share", action="store_true")
+    args = ap.parse_args()
+
+    if args.onnx and os.path.exists(args.onnx):
+        sess, in_name = load_onnx_session(args.onnx)
+        def _fn(img): return predict_onnx(sess, in_name, img, args.img_size)
+        title = "csPWS Segmentation — ONNX (CoreML/CPU auto)"
+    elif os.path.exists(args.ckpt):
+        model, device, img_size = load_torch_model(args.ckpt, 3, args.img_size)
+        def _fn(img): return predict_torch(model, device, img, img_size)
+        title = "csPWS Segmentation — PyTorch"
+    else:
+        raise FileNotFoundError(
+            f"No weights found. Expected either:\n"
+            f" - PyTorch ckpt at {args.ckpt}\n"
+            f" - OR ONNX at {args.onnx or 'exports/model.onnx'}\n"
+        )
+
+    demo = gr.Interface(
+        fn=_fn,
+        inputs=gr.Image(type="pil", label="Upload csPWS composite (Σ/BF/SW)"),
+        outputs=gr.Image(type="pil", label="Predicted Mask"),
+        title=title,
+        allow_flagging="never"
+    )
+    demo.launch(server_port=args.port, share=args.share)
+
+if __name__ == "__main__":
+    main()