app.py

import hashlib
import os
import sys
from io import StringIO
from pathlib import Path

import gradio as gr
import matplotlib.pyplot as plt
import pandas as pd

from usalign_runner import USalignRunner
from utils import calculate_md5, run_community_analysis, run_usalign, save_pdb_files

os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Structure-Based Similarity Network")

    with gr.Row():
        file_input = gr.File(
            label="Upload PDB Files",
            file_count="multiple",
            file_types=[".pdb"],
        )

    output = gr.Textbox(
        label="Upload Results", lines=5, max_lines=5, container=True  # 默认显示行数  # 最大可见行数（超过后自动滚动）
    )

    threshold = gr.Slider(minimum=0, maximum=1, value=0.75, label="Threshold")

    with gr.Row():

        submit_btn = gr.Button("Upload Files")
        run_usalign_btn = gr.Button("Run USalign")
        community_btn = gr.Button("Run Community")

    md5_hash = gr.State("")
    with gr.Tab("USalign Results"):
        results_df = gr.DataFrame(
            label="USalign Results",
            wrap=True,
        )
    with gr.Tab("TM Matrix"):
        # Add new output components for community analysis with height limits
        tm_matrix_output = gr.DataFrame(label="TM Matrix", wrap=True, show_label=True)
    with gr.Tab("Newick Tree"):
        newick_output = gr.Textbox(
            label="Newick Tree", lines=5, max_lines=10, container=True  # 默认显示行数  # 最大可见行数（超过后自动滚动）
        )
    # with gr.Tab("Structure Similarity Network"):
    #     network_plot = gr.Plot(label="Structure Similarity Network")

    # Combine download buttons into a single row
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Download Results")
            download_tm = gr.File(label="Download Files")

    submit_btn.click(fn=save_pdb_files, inputs=[file_input], outputs=output)

    def update_md5_hash(files):
        if files:
            return calculate_md5(files)
        return ""

    file_input.change(fn=update_md5_hash, inputs=[file_input], outputs=[md5_hash])

    run_usalign_btn.click(fn=run_usalign, inputs=[md5_hash], outputs=[results_df])

    def process_community_analysis(results_df, md5_hash, threshold):
        if results_df.empty:
            return None, None, None

        results = run_community_analysis(results_df, "./data", md5_hash, threshold)

        if "Error" in results:
            return None, None, None

        # Prepare download files

        return (
            results["tm_matrix"],
            results["newick_str"],
            # results["network_fig"],
            results["files"],
        )

    community_btn.click(
        fn=process_community_analysis,
        inputs=[results_df, md5_hash, threshold],
        outputs=[
            tm_matrix_output,
            newick_output,
            download_tm,
        ],
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0")