Update app.py

app.py

@@ -1,5 +1,6 @@
 import html
 import re
+import shutil
 import uuid
 from typing import Tuple, List
 from datetime import datetime
@@ -170,13 +171,12 @@ def plot_optimization_metrics(display_df, method="AA2CDS", figsize=(18, 12)):
     """
 
     # 定义要绘制的5个指标及其显示名称
-    metrics = ['GC', 'GC_head', 'CAI', 'CAI_head', 'ENC']
+    metrics = ['GC', 'GC_head', 'CAI', 'CAI_head']
     metric_titles = {
         'GC': 'GC Content',
         'GC_head': "5' GC Content",
         'CAI': 'Codon Adaptation Index (CAI)',
         'CAI_head': "5' CAI",
-        'ENC': 'Effective Number of Codons (ENC)'
     }
 
     # 验证数据列是否存在
@@ -185,7 +185,7 @@ def plot_optimization_metrics(display_df, method="AA2CDS", figsize=(18, 12)):
         raise ValueError(f"DataFrame missing required columns: {missing_cols}")
 
     # 创建2x3的子图布局（最后一个位置留空）
-    fig, axes = plt.subplots(2, 3, figsize=figsize)
+    fig, axes = plt.subplots(2, 2, figsize=figsize)
     axes = axes.flatten()  # 展平为1D数组
 
     # 设置全局样式
@@ -392,7 +392,7 @@ def optimize_cds(protein_seq, species, codon_usage_table, method, status_msg,opt
     status_msg = log(f"   • Protein length: {len(protein_seq)} aa")
 
     timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
-    dirout = f'tmp/{timestamp}_{species}/'
+    dirout = f'tmp1/{timestamp}_{species}/'
     os.makedirs(dirout, exist_ok=True)
 
     task = 'predict_web'
@@ -402,119 +402,131 @@ def optimize_cds(protein_seq, species, codon_usage_table, method, status_msg,opt
     codon_usage_path = f'{dirout}/codon_usage.csv'
     codon_usage_table.to_csv(codon_usage_path, index=False)
 
-    status_msg = log("🔹 Step 2/5: Initial CAI-optimal CDS generation")
-
-    df = pd.DataFrame({'id': [_id], 'RefSeq_aa': [protein_seq]})
-    df.to_csv(dirout + f'{task}/input.csv', index=False)
-
-    reverse_mapping = {
-        "Mus_musculus": "mouse",
-        "Escherichia_coli": "Ec",
-        "Saccharomyces_cerevisiae": "Sac",
-        "Pichia": "Pic",
-        "Homo_sapiens": "Human"
-    }
-    species = reverse_mapping[species]
-    df['species'] = species
+    try:
+        status_msg = log("🔹 Step 2/5: Initial CAI-optimal CDS generation")
+
+        df = pd.DataFrame({'id': [_id], 'RefSeq_aa': [protein_seq]})
+        df.to_csv(dirout + f'{task}/input.csv', index=False)
+
+        reverse_mapping = {
+            "Mus_musculus": "mouse",
+            "Escherichia_coli": "Ec",
+            "Saccharomyces_cerevisiae": "Sac",
+            "Pichia": "Pic",
+            "Homo_sapiens": "Human"
+        }
+        species = reverse_mapping[species]
+        df['species'] = species
+
+        codon_instance = {species: Codon(codon_usage_path, rna=False)}
+        df['cai_best_nn'] = df.apply(
+            lambda x: codon_instance[x['species']].cai_opt_codon(x['RefSeq_aa']), axis=1
+        )
 
-    codon_instance = {species: Codon(codon_usage_path, rna=False)}
-    df['cai_best_nn'] = df.apply(
-        lambda x: codon_instance[x['species']].cai_opt_codon(x['RefSeq_aa']), axis=1
-    )
+        status_msg = log("🔹 Step 3/5: Fragmentation & translation consistency check")
+
+        fragments_list = df.apply(
+            lambda x: process_nucleotide_sequences(
+                x['cai_best_nn'],
+                max_nn_length=1200,
+                step=300,
+                pad_char='_',
+                meta_dict={'_id': x['id'], 'species': x['species']}
+            ),
+            axis=1
+        )
 
-    status_msg = log("🔹 Step 3/5: Fragmentation & translation consistency check")
-
-    fragments_list = df.apply(
-        lambda x: process_nucleotide_sequences(
-            x['cai_best_nn'],
-            max_nn_length=1200,
-            step=300,
-            pad_char='_',
-            meta_dict={'_id': x['id'], 'species': x['species']}
-        ),
-        axis=1
+        expanded_data = pd.DataFrame([item for sublist in fragments_list for item in sublist])
+        expanded_data['truncated_aa'] = expanded_data['truncated_nn'].apply(translate)
+        expanded_data = expanded_data.rename(columns={'truncated_nn': 'cai_best_nn'})
+        expanded_data.to_csv(dirout + f'{task}/TS.csv', index=False)
+
+        status_msg = log("🔹 Step 4/5: Multi-seed neural optimization")
+        # seeds = ['1337', '42', '2022', '2023', '2024', '2025']
+        seeds = optimize_seed.split(',')
+        status_msg = log(f"   • Seeds: {', '.join(seeds)}")
+
+        parser = get_pretraining_args()
+        args = parser.parse_args()
+        args.downstream_data_path = dirout
+        args.task = task
+        args.predict = True
+        args.mlm_pretrained_model_path = 'checkpoint/AA2CDS.pth'
+
+        tmps = []
+        df_trun = pd.read_csv(dirout + f'{task}/TS.csv')
+
+        for seed in seeds:
+            status_msg = log(f"   ⏳ Running inference (seed={seed})")
+            args.seed = seed
+            args.out_dir = f'{dirout}/{seed}'
+            os.makedirs(args.out_dir, exist_ok=True)
+
+            inference(args)
+
+            fpred = f'{args.out_dir}/{task}/TS_pred.csv'
+            os.system(f'cat {fpred}')
+            df_pred = pd.read_csv(fpred)
+            df_info = df_pred.merge(df_trun)
+            print(len(df_info), df_info.columns)
+            seq = assemble_fragments(df_info)
+            analyzer = CodonUsageAnalyzer(codon_usage_path)
+
+            result = single_seq_analysis(seq, _id, codon_usage_path)
+            result.update({
+                'GC': round((seq.count("G") + seq.count("C")) / len(seq), 4),
+                'GC_head': round((seq[:60].count("G") + seq[:60].count("C")) / len(seq[:60]), 4),
+                'CAI': round(analyzer.calculate_CAI(seq), 4),
+                'CAI_head': round(analyzer.calculate_CAI(seq[:60]), 4),
+                '_id': f'seed_{seed}',
+                'CDS_Full': seq,
+                'CDS': seq[:30] + "..." if len(seq) > 30 else seq,
+                'species': species,
+            })
+
+            tmps.append(pd.DataFrame({k: [v] for k, v in result.items()}))
+
+        tmp_df = pd.concat(tmps, ignore_index=True)
+        tmp_df = tmp_df.sort_values(by='CAI', ascending=False)
+        tmp_df.to_csv(f'{dirout}/results.csv', index=False)
+
+        status_msg = log("🔹 Step 5/5: Ranking & visualization")
+
+        display_df = tmp_df[['_id', 'GC', 'GC_head', 'CAI', 'CAI_head', 'ENC', 'species', 'CDS']]
+
+        # 使用示例
+        fig_df = tmp_df[['GC', 'GC_head', 'CAI', 'CAI_head']]  # 只选需要的列
+        fig, axes = plot_optimization_metrics(fig_df, method=method)
+        plt.savefig(f'{dirout}/optimization_metrics.png', dpi=300, bbox_inches='tight')
+        plt.show()
+
+        # fig, ax = plt.subplots(figsize=(10, 6))
+        # scores = display_df["GC"].astype(float).tolist()
+        # bars = ax.bar(range(1, len(scores) + 1), scores, alpha=0.7)
+        # ax.set_xlabel("Sequence Rank")
+        # ax.set_ylabel("GC Content")
+        # ax.set_title(f"CDS Optimization Results ({method})")
+        # ax.grid(True, alpha=0.3)
+
+        # for i in range(min(5, len(bars))):
+        #     bars[i].set_color('orange')
+
+        status_msg = log(f"✅ Successfully generated {len(display_df)} optimized CDS sequences")
+        status_msg = log("🎉 Optimization complete")
+    except Exception as e:
+        status_msg = log(f"❌ Error: {e}")
+        None, None,None, status_msg
+
+    src_dir = "/app/tmp/20251220153157_Mus_musculus/42/predict_web"
+    zip_base = "/app/tmp/predict_web_results"  # 不要加 .zip
+
+    zip_path = shutil.make_archive(
+        base_name=zip_base,
+        format="zip",
+        root_dir=src_dir
     )
 
-    expanded_data = pd.DataFrame([item for sublist in fragments_list for item in sublist])
-    expanded_data['truncated_aa'] = expanded_data['truncated_nn'].apply(translate)
-    expanded_data = expanded_data.rename(columns={'truncated_nn': 'cai_best_nn'})
-    expanded_data.to_csv(dirout + f'{task}/TS.csv', index=False)
-
-    status_msg = log("🔹 Step 4/5: Multi-seed neural optimization")
-    # seeds = ['1337', '42', '2022', '2023', '2024', '2025']
-    seeds = optimize_seed.split(',')
-    status_msg = log(f"   • Seeds: {', '.join(seeds)}")
-
-    parser = get_pretraining_args()
-    args = parser.parse_args()
-    args.downstream_data_path = dirout
-    args.task = task
-    args.predict = True
-    args.mlm_pretrained_model_path = 'checkpoint/AA2CDS.pth'
-
-    tmps = []
-    df_trun = pd.read_csv(dirout + f'{task}/TS.csv')
-
-    for seed in seeds:
-        status_msg = log(f"   ⏳ Running inference (seed={seed})")
-        args.seed = seed
-        args.out_dir = f'{dirout}/{seed}'
-        os.makedirs(args.out_dir, exist_ok=True)
-
-        inference(args)
-
-        fpred = f'{args.out_dir}/{task}/TS_pred.csv'
-        df_pred = pd.read_csv(fpred)
-        df_info = df_pred.merge(df_trun)
-        print(len(df_info),df_info.columns)
-        seq = assemble_fragments(df_info)
-        analyzer = CodonUsageAnalyzer(codon_usage_path)
-
-        result = single_seq_analysis(seq, _id, codon_usage_path)
-        result.update({
-            'GC': round((seq.count("G") + seq.count("C")) / len(seq), 4),
-            'GC_head': round((seq[:60].count("G") + seq[:60].count("C")) / len(seq[:60]), 4),
-            'CAI': round(analyzer.calculate_CAI(seq), 4),
-            'CAI_head': round(analyzer.calculate_CAI(seq[:60]), 4),
-            '_id': f'seed_{seed}',
-            'CDS_Full': seq,
-            'CDS': seq[:30] + "..." if len(seq) > 30 else seq,
-            'species': species,
-        })
-
-        tmps.append(pd.DataFrame({k: [v] for k, v in result.items()}))
-
-    tmp_df = pd.concat(tmps, ignore_index=True)
-    tmp_df = tmp_df.sort_values(by='CAI', ascending=False)
-    tmp_df.to_csv(f'{dirout}/results.csv', index=False)
-
-    status_msg = log("🔹 Step 5/5: Ranking & visualization")
-
-    display_df = tmp_df[['_id',  'GC', 'GC_head', 'CAI', 'CAI_head', 'ENC','species','CDS']]
-
-    # 使用示例
-    fig_df = tmp_df[['GC', 'GC_head', 'CAI', 'CAI_head', 'ENC']]  # 只选需要的列
-    fig, axes = plot_optimization_metrics(fig_df, method=method)
-    plt.savefig(f'{dirout}/optimization_metrics.png', dpi=300, bbox_inches='tight')
-    plt.show()
-
-    # fig, ax = plt.subplots(figsize=(10, 6))
-    # scores = display_df["GC"].astype(float).tolist()
-    # bars = ax.bar(range(1, len(scores) + 1), scores, alpha=0.7)
-    # ax.set_xlabel("Sequence Rank")
-    # ax.set_ylabel("GC Content")
-    # ax.set_title(f"CDS Optimization Results ({method})")
-    # ax.grid(True, alpha=0.3)
-
-
-
-    # for i in range(min(5, len(bars))):
-    #     bars[i].set_color('orange')
-
-    status_msg = log(f"✅ Successfully generated {len(display_df)} optimized CDS sequences")
-    status_msg = log("🎉 Optimization complete")
-
-    return display_df, fig,status_msg
+    return display_df, fig,zip_path,status_msg
 
 def download_cds_results(results_df):
     if results_df is None or len(results_df) == 0:
@@ -1115,41 +1127,43 @@ class MaoTaoWeb:
                 with gr.Row():
                     results_table = gr.Dataframe(
                         label="Optimization Results",
-                        headers=["Rank", "Sequence", "GC%", "tRNA", "Usage", "MFE", "Score"],
-                        datatype=["number", "str", "str", "str", "str", "str", "str"],
-                        col_count=(7, "fixed"),
+                        headers=['_id',  'GC', 'GC_head', 'CAI', 'CAI_head', 'ENC','species','CDS'],
+                        datatype=["str","number", "number", "number", "number", "number", "str", "str"],
+                        col_count=(8, "fixed"),
                         wrap=True
                     )
 
                 optimization_plot = gr.Plot(label="Score Distribution")
 
                 with gr.Row():
-                    download_cds_btn = gr.Button("📥 Download CDS Results", variant="secondary")
-                    cds_download_file = gr.File(label="Download File", visible=False)
+                    # download_cds_btn = gr.Button("📥 Download CDS Results", variant="secondary")
+                    # cds_download_file = gr.File(label="Download File", visible=False)
+
+                    download_btn = gr.DownloadButton(
+                        label="⬇ Download all results (ZIP)",
+                        value='predict_web_results.zip',
+                    )
 
                 def optimize_and_update(protein_seq, species, codon_usage_table,method,optimize_seed):
                     status_msg = f"🔄 Optimizing CDS sequence using {method} method..."
                     # 执行优化
-                    df, plot,status_msg = optimize_cds(protein_seq, species,codon_usage_table, method,status_msg,optimize_seed)
+                    df, plot,zip_path,status_msg = optimize_cds(protein_seq, species,codon_usage_table, method,status_msg,optimize_seed)
                     # 最终状态
 
                     # final_status = f"✅ Optimization complete! Generated {len(df)} sequences with {variants:,} potential variants"
                     # self.status_display.update(final_status)
-                    return df, plot,status_msg
+                    return df, plot,zip_path,status_msg
 
                 optimize_btn.click(
                     optimize_and_update, # protein_seq, species, codon_usage_table,method
                     inputs=[protein_seq, species,codon_usage_table,method,optimize_seed],
-                    outputs=[results_table, optimization_plot, optimize_log]
+                    outputs=[results_table, optimization_plot,download_btn, optimize_log]
                 )
 
                 cds_example_btn.click(lambda: EXAMPLE_PROTEIN, outputs=protein_seq)
 
-                download_cds_btn.click(
-                    download_cds_results,
-                    inputs=results_table,
-                    outputs=cds_download_file
-                )
+
+
 
     def resources_tab(self):
         with gr.Tab("📚 Resources"):