parse rf3d output folder

app.py

@@ -10,6 +10,7 @@ from lightning.fabric import seed_everything
 from rfd3.engine import RFD3InferenceConfig, RFD3InferenceEngine
 from utils import download_weights
 from utils.pipelines import test_rfd3_from_notebook, unconditional_generation
+from gradio_molecule3d import Molecule3D
 
 
 download_weights()
@@ -50,11 +51,61 @@ with gr.Blocks(title="RFD3 Test") as demo:
             maximum=200
         )
 
+    gen_directory = gr.State(None)
+    gen_results = gr.State(None)
     gen_btn = gr.Button("Run Unconditional Generation")
-    gen_output = gr.Textbox(label="Generation Result")
-    gen_btn.click(unconditional_generation, inputs=[num_batches, num_designs_per_batch, length], outputs=gen_output)
+    #gen_output = gr.Textbox(label="Generation Result")
+    gen_btn.click(unconditional_generation, inputs=[num_batches, num_designs_per_batch, length], outputs=[gen_directory, gen_results])
 
 
+    # New visualize section
+    #with gr.Row():
+    #    viz_btn = gr.Button("Visualize", visible=True)
+    #    batch_dropdown = gr.Dropdown(
+    #        choices=[], 
+    #        label="Select Batch",
+    #        visible=False
+    #    )
+    #    design_dropdown = gr.Dropdown(
+    #        choices=[], 
+    #        label="Select Design",
+    #        visible=False
+    #    )
+    #    viewer = Molecule3D(visible=False)
+#
+    #def toggle_visualize(result):
+    #    if result is None:
+    #        return gr.Dropdown(visible=False), gr.Dropdown(visible=False), Molecule3D(visible=False)
+    #    batches = sorted(list({d["batch"] for d in result}))
+    #    return (
+    #        gr.update(choices=batches, visible=True),  # Batch dropdown
+    #        gr.update(choices=[], visible=True),       # Design empty initially
+    #        gr.update(visible=False)
+    #    )
+#
+    #def update_designs(batch, result):
+    #    if batch is None:
+    #        return gr.update(choices=[])
+    #    designs = sorted(list({d["design"] for d in result if d["batch"] == batch}))
+    #    return gr.update(choices=designs)
+#
+    #def load_viewer(batch, design, result):
+    #    if batch is None or design is None:
+    #        return gr.update(visible=False)
+    #    pdb_data = next(d["pdb"] for d in result if d["batch"] == int(batch) and d["design"] == int(design))
+    #    return gr.update(value=pdb_data, visible=True, reps=[{"style": "cartoon"}])  # Customize style
+#
+    ## Events
+    #viz_btn.click(toggle_visualize, inputs=gen_directory, outputs=[batch_dropdown, design_dropdown, viewer])
+    #batch_dropdown.change(update_designs, inputs=[batch_dropdown, gen_directory], outputs=design_dropdown)
+    #batch_dropdown.select(fn=update_designs, inputs=[batch_dropdown, gen_directory], outputs=design_dropdown)  # For selection
+    #gr.Dropdown.select(update_designs, batch_dropdown, design_dropdown).then(  # Chain
+    #    lambda b, d, r: load_viewer(b, d, r),
+    #    inputs=[batch_dropdown, design_dropdown, gen_directory],
+    #    outputs=viewer
+    #)
+
+    
 
 if __name__ == "__main__":
     demo.launch()

utils/pipelines.py

@@ -5,6 +5,9 @@ import time
 import os
 import spaces
 import subprocess
+from Bio.PDB import MMCIFParser, PDBIO
+import gzip
+
 
 
 @spaces.GPU(duration=300)  
@@ -54,10 +57,9 @@ def unconditional_generation(num_batches, num_designs_per_batch, length):
 
     session_hash = gr.Request().session_hash
     time_stamp = time.strftime("%Y-%m-%d-%H-%M-%S")
-    directory = f"./outputs/session_{session_hash}_{time_stamp}"
+    directory = f"./outputs/unconditional_generation/session_{session_hash}_{time_stamp}"
     os.makedirs(directory, exist_ok=False)
 
-
     try: 
         model = RFD3InferenceEngine(**config)
         outputs = model.run(
@@ -66,10 +68,39 @@ def unconditional_generation(num_batches, num_designs_per_batch, length):
             n_batches=num_batches,      # Generate 1 batch
         )
 
-        cmd = f"ls -R {directory}"
+        results = {}
+        for batch in range(num_batches):
+            for design in range(num_designs_per_batch):
+                file_name = os.path.joint(directory, f"_{batch}_{design}.cif.gz")
+                results.append({"batch": batch, "design": design, "file": file_name, "pdb": cif_gz_to_pdb(file_name)})
+            
+        print(results)
+        return directory, results
+    
+    except Exception as e:
+        return f"Error: {str(e)}"
+    
+def collect_outputs(gen_directory, num_batches, num_designs_per_batch):
+    try:
+        cmd = f"ls -R {gen_directory}"
         file_list = subprocess.check_output(cmd, shell=True).decode()
-
         return file_list
-    
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"Error: {str(e)}"
+    
+
+def cif_gz_to_pdb(cif_gz_path):
+    """Convert .cif.gz to PDB string for viewer."""
+    # Decompress & parse
+    parser = MMCIFParser(QUIET=True)
+    with gzip.open(cif_gz_path, 'rt') as f:
+        struct = parser.get_structure('model', f)
+    
+    # Write to string
+    io = PDBIO()
+    io.set_structure(struct)
+    pdb_lines = []
+    class StringIO:
+        def write(self, s): pdb_lines.append(s)
+    io.save(StringIO())
+    return ''.join(pdb_lines)