feat: add bigwig export option

app.py

@@ -3,10 +3,17 @@ import uuid
 import tempfile
 import numpy as np
 import gradio as gr
-import matplotlib.pyplot as plt
 import asyncio
 
+# Set matplotlib to use non-interactive backend before importing pyplot
+# This is required for Gradio which runs on worker threads
+import matplotlib
+matplotlib.use('Agg')
+
+import matplotlib.pyplot as plt
+
 from ntv3_tracks_pipeline import load_ntv3_tracks_pipeline, BED_ELEMENT_COLORS
+from bigwig_export import create_bigwig_zip, _softmax_last
 
 
 # -----------------------------
@@ -54,12 +61,6 @@ load_pipeline(MODEL_ID, DEFAULT_SPECIES)
 # -----------------------------
 # Helpers
 # -----------------------------
-def _softmax_last(x: np.ndarray) -> np.ndarray:
-    x = x - x.max(axis=-1, keepdims=True)
-    ex = np.exp(x)
-    return ex / ex.sum(axis=-1, keepdims=True)
-
-
 def _global_stride(L: int, target: int) -> int:
     if target <= 0 or L <= target:
         return 1
@@ -289,7 +290,7 @@ def predict(
         "plot_target_points": PLOT_TARGET_POINTS,
     }
 
-    return fig, png_path, meta
+    return fig, png_path, meta, out, bigwig_selected, bed_elements
 
 
 # -----------------------------
@@ -691,6 +692,14 @@ with gr.Blocks(title="NTv3 Tracks Demo") as demo:
     
     plot = gr.Plot(label="", elem_id="tracks_plot")
     export_png = gr.File(elem_id="export_png_hidden", interactive=False)
+    
+    # State to store prediction output and selections for BigWig export
+    prediction_state = gr.State(value=None)
+    bigwig_selected_state = gr.State(value=[])
+    bed_elements_state = gr.State(value=[])
+    
+    download_bigwig_btn = gr.Button("📥 Download tracks as BigWig files (ZIP)", variant="secondary")
+    export_bigwig = gr.File(label="Download BigWig files", visible=False)
 
     with gr.Accordion("Meta (click to expand)", open=False):
         meta = gr.JSON(label="Meta")
@@ -750,9 +759,25 @@ with gr.Blocks(title="NTv3 Tracks Demo") as demo:
     btn.click(
         fn=predict,
         inputs=[seq, species, chrom, start, end, use_coords, bigwig_selected, bed_elements],
-        outputs=[plot, export_png, meta],
+        outputs=[plot, export_png, meta, prediction_state, bigwig_selected_state, bed_elements_state],
         api_name="predict",
     )
+    
+    def download_bigwig_zip(out, bw_selected, bed_selected):
+        """Create and return BigWig zip file."""
+        try:
+            zip_path = create_bigwig_zip(out, bw_selected, bed_selected)
+            return gr.update(value=zip_path, visible=True)
+        except ImportError as e:
+            raise gr.Error("pyBigWig is required for BigWig export. Install with: pip install pyBigWig")
+        except Exception as e:
+            raise gr.Error(f"Error creating BigWig files: {str(e)}")
+    
+    download_bigwig_btn.click(
+        fn=download_bigwig_zip,
+        inputs=[prediction_state, bigwig_selected_state, bed_elements_state],
+        outputs=[export_bigwig],
+    )
 
 if __name__ == "__main__":
     demo.launch(

requirements.txt

@@ -5,3 +5,4 @@ gradio>=4.0.0
 pyfaidx
 requests
 matplotlib
+pyBigWig