Spaces:

anton-bushuiev
/

PPIformer

Running on Zero

App Files Files Community

Anton Bushuiev commited on Jan 8, 2024

Commit

dfeea55

1 Parent(s): f230748

Implement "Infinite colors"

Browse files

Files changed (1) hide show

app.py +21 -10

app.py CHANGED Viewed

@@ -1,4 +1,5 @@
 import copy
 import tempfile
 from pathlib import Path
 from functools import partial
@@ -8,6 +9,7 @@ import torch
 from Bio.PDB.Polypeptide import protein_letters_3to1
 from biopandas.pdb import PandasPdb
 from colour import Color
 from mutils.proteins import AMINO_ACID_CODES_1
 from mutils.pdb import download_pdb
@@ -21,6 +23,9 @@ from ppiformer.utils.torch import fill_diagonal
 from ppiformer.definitions import PPIFORMER_WEIGHTS_DIR
 def process_inputs(inputs, temp_dir):
     pdb_code, pdb_path, partners, muts, muts_path = inputs
@@ -110,8 +115,13 @@ def plot_3dmol(pdb_path, ppi_path, muts, attn):
     zoom_atoms = []
     # Cartoon chains
-    colors = [Color(c) for c in ['LimeGreen', 'HotPink', 'RoyalBlue']]
-    chain_to_color = dict(zip(chains, colors))
     for chain in chains:
         styles.append([{"chain": chain}, {"cartoon": {"color": chain_to_color[chain].hex_l, "opacity": 0.6}}])
@@ -136,7 +146,7 @@ def plot_3dmol(pdb_path, ppi_path, muts, attn):
     # Convert style dicts to JS lines
     styles = ''.join(['viewer.addStyle(' + ', '.join([str(s).replace("'", '"') for s in dcts]) + ');\n' for dcts in styles])
-    # Connert zoom atoms to 3DMol.js selection and add labels for mutated residues
     zoom_animation_duration = 500
     sel = '{\"or\": [' + ', '.join(["{\"serial\": " + str(a) + "}" for a in zoom_atoms]) + ']}'
     zoom = 'viewer.zoomTo(' + sel + ',' + f'{zoom_animation_duration});'
@@ -146,7 +156,7 @@ def plot_3dmol(pdb_path, ppi_path, muts, attn):
         label = protein_letters_3to1[row['residue_name']] + row['chain_id'] + str(row['residue_number']) + row['insertion']
         styles += 'viewer.addLabel(' + f"\"{label}\"," + "{fontSize:16, fontColor:\"red\", backgroundOpacity: 0.0}," + sel + ');\n'
-    # Construct 3Dmol.js visualization script in HTML
     html = (
         """<!DOCTYPE html>
         <html>
@@ -207,7 +217,10 @@ def predict(models, temp_dir, *inputs):
     print(ppi_path, muts)
     # Predict
-    ddg, attn = predict_ddg(models, ppi_path, muts, return_attn=True)
     # Create dataframe
     ddg = ddg.detach().numpy().tolist()
@@ -221,9 +234,6 @@ def predict(models, temp_dir, *inputs):
 app = gr.Blocks()
 with app:
-    # print('app.theme.background_fill_primary', app.theme.background_fill_primary, type(app.theme.background_fill_primary))
-    # print('app.theme.background_fill_primary', app.theme.background_fill_primary_dark, type(app.theme.background_fill_primary))
-    # print(app.theme.to_dict())
     # Input GUI
     with gr.Row():
@@ -242,10 +252,11 @@ with app:
     examples = gr.Examples(
         examples=[
             ["1BUI", "A,B,C", "SC16A;FC47A;SC16A,FC47A"],
-            ["1KNE", "A,P", ';'.join([f"TP6{a}" for a in AMINO_ACID_CODES_1])]
         ],
         inputs=[pdb_code, partners, muts],
-        label="Examples (press line to fill)"
     )
     # Predict GUI

 import copy
+import random
 import tempfile
 from pathlib import Path
 from functools import partial
 from Bio.PDB.Polypeptide import protein_letters_3to1
 from biopandas.pdb import PandasPdb
 from colour import Color
+from colour import RGB_TO_COLOR_NAMES
 from mutils.proteins import AMINO_ACID_CODES_1
 from mutils.pdb import download_pdb
 from ppiformer.definitions import PPIFORMER_WEIGHTS_DIR
+random.seed(0)
 def process_inputs(inputs, temp_dir):
     pdb_code, pdb_path, partners, muts, muts_path = inputs
     zoom_atoms = []
     # Cartoon chains
+    preferred_colors = ['LimeGreen', 'HotPink', 'RoyalBlue']
+    all_colors = [c[0] for c in RGB_TO_COLOR_NAMES.values()]
+    all_colors = [c for c in all_colors if c not in preferred_colors + ['Black', 'White']]
+    random.shuffle(all_colors)
+    all_colors = preferred_colors + all_colors
+    all_colors = [Color(c) for c in all_colors]
+    chain_to_color = dict(zip(chains, all_colors))
     for chain in chains:
         styles.append([{"chain": chain}, {"cartoon": {"color": chain_to_color[chain].hex_l, "opacity": 0.6}}])
     # Convert style dicts to JS lines
     styles = ''.join(['viewer.addStyle(' + ', '.join([str(s).replace("'", '"') for s in dcts]) + ');\n' for dcts in styles])
+    # Convert zoom atoms to 3DMol.js selection and add labels for mutated residues
     zoom_animation_duration = 500
     sel = '{\"or\": [' + ', '.join(["{\"serial\": " + str(a) + "}" for a in zoom_atoms]) + ']}'
     zoom = 'viewer.zoomTo(' + sel + ',' + f'{zoom_animation_duration});'
         label = protein_letters_3to1[row['residue_name']] + row['chain_id'] + str(row['residue_number']) + row['insertion']
         styles += 'viewer.addLabel(' + f"\"{label}\"," + "{fontSize:16, fontColor:\"red\", backgroundOpacity: 0.0}," + sel + ');\n'
+    # Construct 3Dmol.js visualization script embedded in HTML
     html = (
         """<!DOCTYPE html>
         <html>
     print(ppi_path, muts)
     # Predict
+    try:
+        ddg, attn = predict_ddg(models, ppi_path, muts, return_attn=True)
+    except:
+        raise gr.Error("Prediction failed. Please double check your inputs.")
     # Create dataframe
     ddg = ddg.detach().numpy().tolist()
 app = gr.Blocks()
 with app:
     # Input GUI
     with gr.Row():
     examples = gr.Examples(
         examples=[
             ["1BUI", "A,B,C", "SC16A;FC47A;SC16A,FC47A"],
+            ["1KNE", "A,P", ';'.join([f"TP6{a}" for a in AMINO_ACID_CODES_1])],
+            ["1C4Z", "A,B,C,D", "FA690A;KD100A"]
         ],
         inputs=[pdb_code, partners, muts],
+        label="Examples (press line to fill inputs)"
     )
     # Predict GUI