Update app.py

app.py

@@ -1,4 +1,5 @@
 import os,sys
+import plotly.graph_objects as go
 
 # install environment goods
 #os.system("pip -q install dgl -f https://data.dgl.ai/wheels/cu113/repo.html")
@@ -54,13 +55,11 @@ args['strand_bias'] = 0.0
 args['loop_bias'] = 0.0
 args['helix_bias'] = 0.0
 
-
-
 def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bias, 
                     secondary_structure, aa_bias, aa_bias_potential, 
-                    #target_charge, target_ph, charge_potential, 
                     num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
                     contigs, pssm, seq_mask, str_mask, rewrite_pdb):
+
     
     dssp_checkpoint = './SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt'
     og_checkpoint = './SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt'
@@ -222,7 +221,9 @@ def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bia
         yield output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)
     
     output_seq, output_pdb, plddt = S.get_outputs()
-    yield output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)
+ 
+    return output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)
+
 
 def get_plddt_plot(plddt_data, max_t):
     x = [i+1 for i in range(len(plddt_data))]
@@ -394,7 +395,19 @@ def get_pssm(fasta_msa, input_pssm):
     return fig, outdir
 
 
-#toggle options
+
+
+# 히어로 능력치 계산 함수 추가
+def calculate_hero_stats(helix_bias, strand_bias, loop_bias, hydrophobic_score):
+    stats = {
+        'strength': strand_bias * 20,  # 베타시트 구조 기반
+        'flexibility': helix_bias * 20, # 알파헬릭스 구조 기반
+        'speed': loop_bias * 5,        # 루프 구조 기반
+        'defense': abs(hydrophobic_score) if hydrophobic_score else 0
+    }
+    return stats
+
+##toggle options
 def toggle_seq_input(choice):
     if choice == "protein length":
         return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
@@ -407,16 +420,116 @@ def toggle_secondary_structure(choice):
     elif choice == "explicit":
         return gr.update(visible=False, value=None),gr.update(visible=False, value=None),gr.update(visible=False, value=None),gr.update(visible=True, value=None)
 
+def generate_hero(name, strength, flexibility, speed, defense, size, abilities):
+    try:
+        # protein_diffusion_model 호출
+        results = protein_diffusion_model(
+            sequence=None,
+            seq_len=size,
+            helix_bias=flexibility,
+            strand_bias=strength,
+            loop_bias=speed,
+            secondary_structure=None,
+            aa_bias=None,
+            aa_bias_potential=None,
+            num_steps="25",
+            noise="normal",
+            hydrophobic_target_score=str(-defense),
+            hydrophobic_potential="2",
+            contigs=None,
+            pssm=None,
+            seq_mask=None,
+            str_mask=None,
+            rewrite_pdb=None
+        )
+        
+        # 능력치 계산
+        stats = calculate_hero_stats(flexibility, strength, speed, defense)
+        
+        # 결과 반환
+        return (
+            results[2],  # 3D 구조
+            create_radar_chart(stats),  # 능력치 차트
+            generate_hero_description(name, stats, abilities)  # 히어로 설명
+        )
+    except Exception as e:
+        return (
+            gr.HTML("에러가 발생했습니다"),
+            None,
+            f"에러: {str(e)}"
+        )
+
+def create_radar_chart(stats):
+    # 레이더 차트 생성 로직
+    categories = list(stats.keys())
+    values = list(stats.values())
+    
+    fig = go.Figure(data=go.Scatterpolar(
+        r=values,
+        theta=categories,
+        fill='toself'
+    ))
+    
+    fig.update_layout(
+        polar=dict(
+            radialaxis=dict(
+                visible=True,
+                range=[0, 1]
+            )),
+        showlegend=False
+    )
+    
+    return fig
 
-# Define the Gradio interface
-with gr.Blocks(theme='ParityError/Interstellar') as demo:
+def generate_hero_description(name, stats, abilities):
+    # 히어로 설명 생성 로직
+    description = f"""
+    히어로 이름: {name}
     
- 
+    주요 능력:
+    - 근력: {'★' * int(stats['strength'] * 5)}
+    - 유연성: {'★' * int(stats['flexibility'] * 5)}
+    - 스피드: {'★' * int(stats['speed'] * 5)}
+    - 방어력: {'★' * int(stats['defense'] * 5)}
+    
+    특수 능력: {', '.join(abilities)}
+    """
+    return description
 
-    #with gr.Row().style(equal_height=False):
+# Gradio 인터페이스 수정
+with gr.Blocks(theme='ParityError/Interstellar') as demo:
     with gr.Row():
         with gr.Column():
+            gr.Markdown("# 🦸‍♂️ 슈퍼히어로 단백질 만들기")
+            
             with gr.Tabs():
+                with gr.TabItem("히어로 디자인"):
+                    # 히어로 기본 정보
+                    hero_name = gr.Textbox(label="히어로 이름", placeholder="당신의 히어로 이름을 지어주세요!")
+                    
+                    # 능력치 설정
+                    gr.Markdown("### 💪 히어로 능력치 설정")
+                    with gr.Row():
+                        strength = gr.Slider(minimum=0.0, maximum=0.05, label="💪 초강력(근력)", value=0.02)
+                        flexibility = gr.Slider(minimum=0.0, maximum=0.05, label="🤸‍♂️ 유연성", value=0.02)
+                    
+                    with gr.Row():
+                        speed = gr.Slider(minimum=0.0, maximum=0.20, label="⚡ 스피드", value=0.1)
+                        defense = gr.Slider(minimum=-10, maximum=10, label="🛡️ 방어력", value=0)
+                    
+                    # 히어로 크기 설정
+                    hero_size = gr.Slider(minimum=50, maximum=200, label="히어로 크기", value=100)
+                    
+                    # 특수 능력 설정
+                    with gr.Accordion("🌟 특수 능력", open=False):
+                        special_ability = gr.CheckboxGroup(
+                            choices=["자가 회복", "원거리 공격", "방어막 생성"],
+                            label="특수 능력 선택"
+                        )
+                    
+                    # 생성 버튼
+                    create_btn = gr.Button("히어로 생성!", variant="primary")
+
                 with gr.TabItem("Inputs"):
                     gr.Markdown("""## INPUTS""")
                     gr.Markdown("""#### Start Sequence
@@ -453,15 +566,6 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                         with gr.Row():
                             aa_bias = gr.Textbox(label="aa bias", lines=1, placeholder='specify one letter AA and fraction to bias, for example W0.1 or M0.1,K0.1' )
                             aa_bias_potential = gr.Textbox(label="aa bias scale", lines=1, placeholder='AA Bias potential scale (recomended range 1.0-5.0)')
-                    
-                    '''
-                    with gr.Accordion(label='Charge Bias',open=False):
-                        gr.Markdown("""Bias for a specified net charge at a particular pH using the boxes below""")
-                        with gr.Row():
-                            target_charge = gr.Textbox(label="net charge", lines=1, placeholder='net charge to target')
-                            target_ph = gr.Textbox(label="pH", lines=1, placeholder='pH at which net charge is desired')
-                            charge_potential = gr.Textbox(label="charge potential scale", lines=1, placeholder='charge potential scale (recomended range 1.0-5.0)')
-                    '''
 
                     with gr.Accordion(label='Hydrophobic Bias',open=False):
                         gr.Markdown("""Bias for or against hydrophobic composition, to get more soluble proteins, bias away with a negative target score (ex. -5)""")
@@ -473,10 +577,9 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                         gr.Markdown("""Increasing T to more steps can be helpful for harder design challenges, sampling from different distributions can change the sequence and structural composition""")
                         with gr.Row():
                             num_steps = gr.Textbox(label="T", lines=1, placeholder='number of diffusion steps (25 or less will speed things up)')
-                            noise = gr.Dropdown(['normal','gmm2 [-1,1]','gmm3 [-1,0,1]'], label='noise type', value='normal') 
-                
-                with gr.TabItem("Motif Selection"):
+                            noise = gr.Dropdown(['normal','gmm2 [-1,1]','gmm3 [-1,0,1]'], label='noise type', value='normal')
 
+                with gr.TabItem("Motif Selection"):
                     gr.Markdown("""### Motif Selection Preview""")
                     gr.Markdown('Contigs explained: to grab residues (seq and str) on a pdb chain you will provide the chain letter followed by a range of residues as indexed in the pdb file for example (A3-10) is the syntax to select residues 3-10 on chain A (the chain always needs to be specified). To add diffused residues to either side of this motif you can specify a range or discrete value without a chain letter infront. To add 15 residues before the motif and 20-30 residues (randomly sampled) after use the following syntax: 15,A3-10,20-30 commas are used to separate regions selected from the pdb and designed (diffused) resiudes which will be added. ')
                     pdb_id_code = gr.Textbox(label="PDB ID", lines=1, placeholder='INPUT PDB ID TO FETCH (ex. 1DPX)', visible=True)
@@ -499,11 +602,21 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                     pssm_view = gr.Plot(label='PSSM Viewer')
                     pssm_gen_btn = gr.Button("Generate PSSM")
 
-
-            btn = gr.Button("GENERATE")
-
-        #with gr.Row():
         with gr.Column():
+            gr.Markdown("## 🦸‍♂️ 히어로 프로필")
+            
+            # 히어로 시각화
+            hero_structure = gr.HTML(label="히어로 구조")
+            
+            # 능력치 레이더 차트
+            hero_stats = gr.Plot(label="능력치 분석")
+            
+            # 히어로 설명
+            hero_description = gr.Textbox(label="히어로 특성", lines=3)
+            
+            # 다운로드 버튼
+            download_btn = gr.Button("히어로 데이터 다운로드")
+
             gr.Markdown("""## OUTPUTS""")
             gr.Markdown("""#### Confidence score for generated structure at each timestep""")
             plddt_plot = gr.Plot(label='plddt at step t')
@@ -513,53 +626,17 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
             output_pdb = gr.File(label="PDB file")
             gr.Markdown("""#### Structure viewer""")
             output_viewer = gr.HTML()
-    '''
-    gr.Markdown("""### Don't know where to get started? Click on an example below to try it out!""")
-    gr.Examples(
-        [["","125",0.0,0.0,0.2,"","","","20","normal",'','','',None,'','',None],
-         ["","100",0.0,0.0,0.0,"","W0.2","2","20","normal",'','','',None,'','',None],
-#         ["","100",0.0,0.0,0.0,
-#        "XXHHHHHHHHHXXXXXXXHHHHHHHHHXXXXXXXHHHHHHHHXXXXSSSSSSSSSSSXXXXXXXXSSSSSSSSSSSSXXXXXXXSSSSSSSSSXXXXXXX",
-#                "","","25","normal",'','','',None,'','',None],
-#         ["XXXXXXXXXXXXXXXXXXXXXXXXXIPDXXXXXXXXXXXXXXXXXXXXXXPEPSEQXXXXXXXXXXXXXXXXXXXXXXXXXXIPDXXXXXXXXXXXXXXXXXXX",
-#          "",0.0,0.0,0.0,"","","","25","normal",'','','',None,'','',None],
-#         ["","",0.0,0.0,0.0,"","","","25","normal",'','',
-#          '9,D10-11,8,D20-20,4,D25-35,65,D101-101,2,D104-105,8,D114-116,15,D132-138,6,D145-145,2,D148-148,12,D161-161,3',
-#          './tmp/PSSM_lysozyme.csv',
-#          'D25-25,D27-31,D33-35,D132-137',
-#          'D26-26','./tmp/150l.pdb']
-         ],
-        inputs=[sequence, 
-                seq_len, 
-                helix_bias, 
-                strand_bias, 
-                loop_bias, 
-                secondary_structure, 
-                aa_bias, 
-                aa_bias_potential,
-                #target_charge, 
-                #target_ph, 
-                #charge_potential, 
-                num_steps, 
-                noise, 
-                hydrophobic_target_score, 
-                hydrophobic_potential,
-                contigs,
-                pssm,
-                seq_mask,
-                str_mask,
-                rewrite_pdb],
-        outputs=[output_seq,
-                 output_pdb,
-                 output_viewer,
-                 plddt_plot],
-        fn=protein_diffusion_model,
-        ) 
-    '''
-    preview_btn.click(get_motif_preview,[pdb_id_code, contigs],[preview_viewer, rewrite_pdb])
 
+    # 이벤트 연결
+    preview_btn.click(get_motif_preview,[pdb_id_code, contigs],[preview_viewer, rewrite_pdb])
     pssm_gen_btn.click(get_pssm,[fasta_msa,input_pssm],[pssm_view, pssm])
+    create_btn.click(
+        generate_hero,
+        inputs=[hero_name, strength, flexibility, speed, defense, hero_size, special_ability],
+        outputs=[hero_structure, hero_stats, hero_description]
+    )
 
+    btn = gr.Button("GENERATE")
     btn.click(protein_diffusion_model, 
                 [sequence, 
                  seq_len, 
@@ -569,9 +646,6 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                  secondary_structure, 
                  aa_bias, 
                  aa_bias_potential, 
-                 #target_charge, 
-                 #target_ph, 
-                 #charge_potential, 
                  num_steps, 
                  noise, 
                  hydrophobic_target_score, 
@@ -587,7 +661,4 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                  plddt_plot])
 
 demo.queue()
-demo.launch(debug=True)
-
-
-
+demo.launch(debug=True)