add visualization with Molecule3D

app.py

@@ -59,53 +59,52 @@ with gr.Blocks(title="RFD3 Test") as demo:
 
 
     # 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
-    #)
+    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()