Switch to gradio UI

app.py

@@ -1,35 +0,0 @@
-import streamlit as st
-
-# Create the Streamlit app
-st.title("DiffDock Pocket: Protein-Ligand Interaction Calculator")
-
-# Create a form block
-with st.form(key="calculation_form"):
-    # File upload for Protein PDB
-    protein_pdb = st.file_uploader("Protein PDB", type=["pdb"])
-
-    # File upload for Ligand SDF
-    ligand_sdf = st.file_uploader("Ligand SDF", type=["sdf"])
-
-    # Numeric input for "Samples per Complex"
-    samples_per_complex = st.number_input("Samples per Complex", min_value=1, max_value=100, value=4, step=1)
-
-    # Boolean checkbox for "Keep Local Structures"
-    keep_local_structures = st.checkbox("Keep Local Structures", value=True)
-
-    # Boolean checkbox for "Save Visualization"
-    save_visualization = st.checkbox("Save Visualization", value=True)
-
-    # Submit button
-    submit_button = st.form_submit_button("Calculate")
-
-    if submit_button:
-        # Implement your calculation logic here
-        st.write("TODO Calculating... (Add your implementation logic)")
-
-        # For demonstration purposes, display the user inputs
-        st.write(f"Protein PDB File: {protein_pdb}")
-        st.write(f"Ligand SDF File: {ligand_sdf}")
-        st.write(f"Samples per Complex: {samples_per_complex}")
-        st.write(f"Keep Local Structures: {keep_local_structures}")
-        st.write(f"Save Visualization: {save_visualization}")

main.py

@@ -8,7 +8,7 @@ response = requests.get("https://git.io/JJkYN")
 labels = response.text.split("\n")
 
 
-def predict(inp):
+def predict(inp, *args, **kwargs):
     inp = transforms.ToTensor()(inp).unsqueeze(0)
     with torch.no_grad():
         prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
@@ -16,15 +16,29 @@ def predict(inp):
     return confidences
 
 
+def calculate(*args, **kwargs) -> str:
+    output_file_path = "main_output.txt"
+    with open(output_file_path, "w") as fi:
+        fi.write(f"args: {args}\n")
+        fi.write(f"kwargs: {kwargs}\n")
+    return output_file_path
+
+
 def run():
-    demo = gr.Interface(
-        fn=predict,
-        inputs=gr.Image(type="pil"),
-        outputs=gr.Label(num_top_classes=3),
+    iface = gr.Interface(
+        fn=calculate,
+        inputs=[
+            gr.File(label="Protein PDB", file_types=[".pdb"]),
+            gr.File(label="Ligand SDF", file_types=[".sdf"]),
+            gr.Number(label="Samples Per Complex", value=4, minimum=1, maximum=100, precision=0),
+            gr.Checkbox(label="Keep Local Structures", value=True),
+            gr.Checkbox(label="Save Visualization", value=True)
+        ],
+        outputs=gr.File(label="Result")
     )
 
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    iface.launch(server_name="0.0.0.0", server_port=7860)
 
 
 if __name__ == "__main__":
-    run()
+    run()