gradio-app / gradio

Build and share delightful machine learning apps, all in Python. 🌟 Star to support our work!
http://www.gradio.app
Apache License 2.0
34.11k stars 2.59k forks source link

gradio-molecule3d can not update representations. #10008

Closed linuxonly801 closed 1 day ago

linuxonly801 commented 1 day ago

Describe the bug

I am working on antibody structure, I would like to color CDR and FR regions of the antibody by different colors.

gradio-molecule3d is used to display the structure. but when I want to change the color of the structure, it successed at the first time when I dropdown select a color, but fails for the rest attempts.

Screencast from 20241121.webm

the attached file is the pdb file I used, please remove the .txt extension before use. antibody1.pdb.txt

Have you searched existing issues? πŸ”Ž

Reproduction

# Import the gradio library
import gradio as gr
from gradio_molecule3d import Molecule3D
from Bio import PDB
from Bio.SeqUtils import seq1 as biopython_seq1
from abnumber import Chain

def parse_pdb_info(pdb_path):
    """_summary_: Parse the pdb file, extract the chain id, length, sequence, and sequence index

    Args:
        pdb_path (str): Path to the pdb file

    Returns:
        list: [[chain_id, seq_str, sequence_index],[chain_id, seq_str, sequence_index]...]
    """
    # Read the PDB file
    parser = PDB.PDBParser(QUIET=True)
    structure = parser.get_structure("protein", pdb_path)

    # Initialize the list to save information
    chain_info = []

    # Iterate through all models, chains, and residues
    for model in structure:
        for chain in model:
            chain_id = chain.id
            sequence = []
            sequence_index = []
            for residue in chain:
                # Check if it is a standard amino acid
                if PDB.is_aa(residue):
                    sequence.append(biopython_seq1(residue.resname))
                    sequence_index.append(residue.full_id[3][1])
            # Extract the sequence and length of the chain
            seq_str = ''.join(sequence)
            chain_info.append((chain_id, seq_str, sequence_index))
    return chain_info

def generate_reps(chain_info, scheme='imgt', cdr_color='redCarbon', fr_color='greenCarbon', show_sidechain=False, sidechain_color='grayCarbon'):
    """Generate reps based on chain information

    Args:
        chain_info (list): [[chain_id, seq_str, sequence_index],[chain_id, seq_str, sequence_index]...]

    Returns:
        list: reps
    """
    reps = []
    # Iterate through the information of each chain
    for chain_id, seq_str, sequence_index in chain_info:
        # Extract the CDR and FR sequences
        chain = Chain(seq_str, scheme=scheme)
        fr1 = chain.fr1_seq
        cdr1 = chain.cdr1_seq
        fr2 = chain.fr2_seq
        cdr2 = chain.cdr2_seq
        fr3 = chain.fr3_seq
        cdr3 = chain.cdr3_seq
        fr4 = chain.fr4_seq
        # Generate reps information
        cdr_fr_list = [fr1, cdr1, fr2, cdr2, fr3, cdr3, fr4]
        for i in range(len(cdr_fr_list)):
            # Display CDR and FR in different colors
            if i % 2 == 0:
                color = fr_color  # It has been tested that the color must be added with "Carbon", otherwise it cannot be displayed
            else:
                color = cdr_color  # It has been tested that the color must be added with "Carbon", otherwise it cannot be displayed

            # Calculate the start and end positions of CDR and FR
            end = sum([len(cdr_fr_list[j]) for j in range(i+1)])
            start = end - len(cdr_fr_list[i]) + 1
            # Convert to sequence index in pdb file
            end_index = sequence_index[end-1]
            start_index = sequence_index[start-1]

            reps.append({
                "model": 0,  # The nth model loaded. It needs to be specified when multiple models are loaded.
                "chain": chain_id,  # Specify the chain to be displayed. If it is empty, all chains are displayed.
                "style": "cartoon",  # Specify the display style. It can be "cartoon", "stick", "surface", "sphere"
                "color": color,
                "residue_range": f"{start_index}-{end_index}",  # Specify the range of residues to be displayed. If it is empty, all residues are displayed.
                # "residue_range": f"{start}-{end}",  # Specify the range of residues to be displayed. If it is empty, all residues are displayed.
            })

        # If the side chain needs to be displayed, add the reps information of the side chain
        if show_sidechain:
            reps.append({
                "model": 0,  # The nth model loaded. It needs to be specified when multiple models are loaded.
                "chain": "",  # Specify the chain to be displayed. If it is empty, all chains are displayed.
                "style": "stick",  # Specify the display style. It can be "cartoon", "stick", "surface", "sphere"
                "color": sidechain_color,
                "residue_range": "",  # Specify the range of residues to be displayed. If it is empty, all residues are displayed.
            })
    return reps

def update_ab_color(molecule_window, scheme, cdr_color, fr_color, show_sidechain, sidechain_color, background_color):
    chain_info = parse_pdb_info(molecule_window.name)
    repsentation = generate_reps(chain_info, scheme=scheme, cdr_color=cdr_color, fr_color=fr_color, show_sidechain=show_sidechain, sidechain_color=sidechain_color)
    return gr.update(value=molecule_window.name, reps=repsentation, label='updated molecule', config={"backgroundColor": background_color}), gr.update(value=repsentation)

def show_sidechain_color_selector(show_sidechain):
    if show_sidechain:
        return gr.update(visible=True, value='grayCarbon')
    else:
        return gr.update(visible=False, value='blackCarbon')

if __name__ == '__main__':
    repsent =    [
    {
        "model": 0,
        "style": "cartoon",
        "color": "orangeCarbon",
    }
    ]

    color_choices = ['redCarbon', 'greenCarbon', 'orangeCarbon', 'blackCarbon', 'blueCarbon', 'grayCarbon', 'cyanCarbon']

    with gr.Blocks(title="Antibody 3D") as demo:
        gr.Markdown("# Molecule3D")
        with gr.Row():
            with gr.Column():
                molecule_window = Molecule3D(label="Molecule3D", reps=repsent, config={"backgroundColor": 'lightgray'})
            with gr.Column():
                number_scheme = gr.Radio(label="Naming scheme", choices=["imgt", "kabat"], value="imgt")
                with gr.Row():
                    cdr_color = gr.Dropdown(choices=color_choices, label="CDR color", value='redCarbon')
                    fr_color = gr.Dropdown(choices=color_choices, label="FR color", value='greenCarbon')
                with gr.Row():
                    show_sidechain = gr.Checkbox(label="Show side chain", value=False)
                    side_chain_color = gr.Dropdown(choices=color_choices, label="Side chain color", value='blackCarbon', visible=False)
                bg_color = gr.Radio(choices=['lightgray', 'black', 'white', 'orange', 'gray'], label="Background color", value='lightgray')
                reps_text = gr.Textbox(label="reps", value='this is the default value')

        # Callback function for changing display colors
        number_scheme.change(fn=update_ab_color, inputs=[molecule_window, number_scheme, cdr_color, fr_color, show_sidechain, side_chain_color, bg_color], outputs=[molecule_window, reps_text])
        cdr_color.change(fn=update_ab_color, inputs=[molecule_window, number_scheme, cdr_color, fr_color, show_sidechain, side_chain_color, bg_color], outputs=[molecule_window, reps_text])
        fr_color.change(fn=update_ab_color, inputs=[molecule_window, number_scheme, cdr_color, fr_color, show_sidechain, side_chain_color, bg_color], outputs=[molecule_window, reps_text])
        show_sidechain.change(fn=show_sidechain_color_selector, inputs=show_sidechain, outputs=side_chain_color)
        side_chain_color.change(fn=update_ab_color, inputs=[molecule_window, number_scheme, cdr_color, fr_color, show_sidechain, side_chain_color, bg_color], outputs=[molecule_window, reps_text])
        bg_color.change(fn=update_ab_color, inputs=[molecule_window, number_scheme, cdr_color, fr_color, show_sidechain, side_chain_color, bg_color ], outputs=[molecule_window, reps_text])
    demo.launch()

Screenshot

No response

Logs

No response

System Info

OS: ubuntu 20.04.6
gradio-molecule3d: 0.0.6

Gradio Environment Information:
------------------------------
Operating System: Linux
gradio version: 5.6.0
gradio_client version: 1.4.3

------------------------------------------------
gradio dependencies in your environment:

aiofiles: 23.2.1
anyio: 4.6.2.post1
audioop-lts is not installed.
fastapi: 0.115.5
ffmpy: 0.4.0
gradio-client==1.4.3 is not installed.
httpx: 0.27.2
huggingface-hub: 0.26.2
jinja2: 3.1.4
markupsafe: 2.1.5
numpy: 2.1.3
orjson: 3.10.11
packaging: 24.2
pandas: 2.2.3
pillow: 11.0.0
pydantic: 2.10.0
pydub: 0.25.1
python-multipart==0.0.12 is not installed.
pyyaml: 6.0.2
ruff: 0.7.4
safehttpx: 0.1.1
semantic-version: 2.10.0
starlette: 0.41.3
tomlkit==0.12.0 is not installed.
typer: 0.13.1
typing-extensions: 4.12.2
urllib3: 2.2.3
uvicorn: 0.32.1
authlib; extra == 'oauth' is not installed.
itsdangerous; extra == 'oauth' is not installed.

gradio_client dependencies in your environment:

fsspec: 2024.10.0
httpx: 0.27.2
huggingface-hub: 0.26.2
packaging: 24.2
typing-extensions: 4.12.2
websockets: 12.0

Severity

Blocking usage of gradio

abidlabs commented 1 day ago

Hi @linuxonly801, as this is a custom component and not a core Gradio component, please open a discussion here: https://huggingface.co/spaces/simonduerr/gradio_molecule3d

cc @duerrsimon