Update app.py

app.py

@@ -21,7 +21,7 @@ import moleculekit
 print(moleculekit.__version__)
 
 
-def update(inp, file, mode):
+def update(inp, file, mode, custom_resids, clustering_threshold):
     try:
         pdb_file = file.name
     except:
@@ -54,9 +54,11 @@ def update(inp, file, mode):
         ids = get_all_protein_resids(
             f"files/{pdb_file}.pdb",
         )
+    elif len(custom_resids)!=0:
+        ids = [int(x) for x in custom_resids.split(',')]
     else:
         ids = get_all_metalbinding_resids(f"files/{pdb_file}.pdb")
-
+    
     voxels, prot_centers, prot_N, prots = processStructures(pdb_file, ids)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     voxels.to(device)
@@ -89,7 +91,7 @@ def update(inp, file, mode):
         writeprobes=True,
         probefile=f"output/probes_{pdb_file}.pdb",
         threshold=7,
-        p=0.15,
+        p=clustering_threshold,
     )
 
     return message, molecule(
@@ -231,7 +233,7 @@ metal3d = gr.Blocks()
 with metal3d:
     gr.Markdown("# Metal3D")
     with gr.Tabs():
-        with gr.TabItem("App"):
+        with gr.TabItem("Input"):
             inp = gr.Textbox(            placeholder="PDB Code or Uniprot identifier or upload file below", label="Input molecule"
             )
             file = gr.File(file_count="single", type="file")
@@ -243,7 +245,9 @@ with metal3d:
                     label="Residues to use for prediction",
                 )
                 custom_resids = gr.Textbox(placeholder="Comma separated list of residues", label="Custom residues")
-            gr.Slider(minimum=0.15,maximum=1, value=0.15, label="Clustering threshold")
+            with gr.Row():
+                clustering_threshold = gr.Slider(minimum=0.15,maximum=1, value=0.15,step=0.05, label="Clustering threshold")
+                distance_cutoff = gr.Slider(minimum=1,maximum=10, value=7,step=1, label="Clustering distance cutoff")
         btn = gr.Button("Run")
     gr.Markdown(
         """ <small>Inference using CPU-only, can be quite slow for more than 20 residues. Use Colab notebook for GPU acceleration</small>
@@ -254,7 +258,7 @@ with metal3d:
     gr.Markdown("# Output")
     out = gr.Textbox(label="status")
     mol = gr.HTML()
-    btn.click(fn=update, inputs=[inp, file, mode], outputs=[out, mol])
+    btn.click(fn=update, inputs=[inp, file, mode, custom_resids, clustering_threshold], outputs=[out, mol])
 
 metal3d.launch()