Update app.py

app.py

@@ -1,177 +1,120 @@
-# -*- coding: utf-8 -*-
+# app.py
 import gradio as gr
 import torch
 import pickle
 import subprocess
 import pandas as pd
-from predictor import GenePredictor
-from tensorflow.keras.models import load_model
-import ml_simplified_tree  # Import the module
-import logging
 import os
 import re
+import logging
+from predictor import GenePredictor
+from tensorflow.keras.models import load_model
+import ml_simplified_tree
 
-# Configure logging
+# --- Logging ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# --------- Load Models (Optional) ---------
+# --- Paths ---
+model_dir = "models"
+boundary_path = os.path.join(model_dir, "best_boundary_aware_model.pth")
+keras_path = os.path.join(model_dir, "best_model.keras")
+kmer_path = os.path.join(model_dir, "kmer_to_index.pkl")
+csv_path = "f gene clean dataset.csv"
+
+# --- Load Models ---
 boundary_model = None
 keras_model = None
 kmer_to_index = None
-model_dir = "models/"  # Adjust this if the models are in a different subdirectory
-try:
-    boundary_path = os.path.join(model_dir, "best_boundary_aware_model.pth")
-    if os.path.exists(boundary_path):
-        boundary_model = GenePredictor(boundary_path)
-        logging.info(f"Boundary model loaded successfully from {boundary_path}")
-    else:
-        logging.warning(f"Boundary model not found at {boundary_path}. Skipping Boundary-Aware step.")
-except Exception as e:
-    logging.error(f"Error loading Boundary model from {boundary_path}: {e}. Skipping Boundary-Aware step.")
-try:
-    keras_path = os.path.join(model_dir, "best_model.keras")
-    kmer_path = os.path.join(model_dir, "kmer_to_index.pkl")
-    if os.path.exists(keras_path) and os.path.exists(kmer_path):
-        keras_model = load_model(keras_path)
-        with open(kmer_path, "rb") as f:
-            kmer_to_index = pickle.load(f)
-        logging.info(f"Keras model and kmer_to_index loaded successfully from {keras_path} and {kmer_path}")
-    else:
-        logging.warning(f"Keras model or kmer_to_index not found at {keras_path} or {kmer_path}. Skipping Keras step.")
-except Exception as e:
-    logging.error(f"Error loading Keras model or kmer_to_index from {keras_path} or {kmer_path}: {e}. Skipping Keras step.")
 
-# --------- Utilities ---------
+if os.path.exists(boundary_path):
+    boundary_model = GenePredictor(boundary_path)
+    logging.info("Boundary model loaded.")
+if os.path.exists(keras_path) and os.path.exists(kmer_path):
+    keras_model = load_model(keras_path)
+    with open(kmer_path, "rb") as f:
+        kmer_to_index = pickle.load(f)
+    logging.info("Keras model and k-mer index loaded.")
+
+# --- Keras Prediction ---
 def predict_with_keras(sequence):
-    global keras_model, kmer_to_index
-    if keras_model is None or kmer_to_index is None:
-        return f"Skipped (Models not loaded): {sequence}"
+    kmers = [sequence[i:i+6] for i in range(len(sequence)-5)]
+    indices = [kmer_to_index.get(kmer, 0) for kmer in kmers]
+    input_arr = torch.tensor([indices])
+    prediction = keras_model.predict(input_arr, verbose=0)[0]
+    return ''.join(str(round(p, 3)) for p in prediction)
+
+# --- Full Pipeline ---
+def run_pipeline(dna_input):
+    dna_input = dna_input.upper()
+    if not re.match('^[ACTGN]+$', dna_input):
+        dna_input = ''.join(c if c in 'ACTGN' else 'N' for c in dna_input)
+
+    # Step 1: Boundary Prediction
+    if boundary_model:
+        predictions, probs, confidence = boundary_model.predict(dna_input)
+        regions = boundary_model.extract_gene_regions(predictions, dna_input)
+        step1_out = regions[0]["sequence"] if regions else dna_input
+    else:
+        step1_out = dna_input
+
+    # Step 2: Keras Prediction
+    step2_out = predict_with_keras(step1_out) if keras_model and kmer_to_index else step1_out
+
+    # Save to FASTA for MAFFT/IQTREE
+    fasta_file = "input_sequence.fasta"
+    with open(fasta_file, "w") as f:
+        f.write(">query\n" + step2_out + "\n")
+
+    # Run MAFFT
+    aligned_file = "aligned.fasta"
     try:
-        kmers = [sequence[i:i+6] for i in range(len(sequence)-5)]
-        indices = [kmer_to_index.get(kmer, 0) for kmer in kmers]
-        input_arr = torch.tensor([indices])
-        prediction = keras_model.predict(input_arr, verbose=0)[0]
-        return "".join(str(round(p, 3)) for p in prediction)
+        subprocess.run(["mafft", "--auto", fasta_file], stdout=open(aligned_file, "w"), check=True)
     except Exception as e:
-        logging.error(f"Error in Keras prediction: {e}")
-        return f"Error in Keras prediction: {e}"
+        aligned_file = None
+        logging.error(f"MAFFT failed: {e}")
 
-def run_full_pipeline(dna_input):
+    # Run IQ-TREE
+    phy_file = "input_sequence.phy"
     try:
-        # Validate input
-        dna_input = dna_input.upper()
-        if not re.match('^[ACTGN]+$', dna_input):
-            logging.warning("Invalid DNA sequence. Replacing non-[ACTGN] with 'N'.")
-            dna_input = ''.join(c if c in 'ACTGN' else 'N' for c in dna_input)
-
-        # Step 1: Boundary-Aware Prediction (Optional)
-        step1_out = dna_input  # Fallback to input if model missing
-        if boundary_model:
-            try:
-                predictions, probs, confidence = boundary_model.predict(dna_input)
-                gene_regions = boundary_model.extract_gene_regions(predictions, dna_input)
-                step1_out = gene_regions[0]["sequence"] if gene_regions else dna_input
-                logging.info(f"Boundary model output: {step1_out[:50]}... (truncated)")
-            except Exception as e:
-                logging.error(f"Boundary model prediction failed: {e}")
-                step1_out = dna_input
-        else:
-            logging.info("Boundary model skipped due to missing or failed load")
-
-        # Step 2: Keras Prediction (Optional)
-        step2_out = step1_out  # Fallback to step1_out if model missing
-        if keras_model and kmer_to_index:
-            try:
-                step2_out = predict_with_keras(step1_out)
-                logging.info(f"Keras model output: {step2_out[:50]}... (truncated)")
-            except Exception as e:
-                logging.error(f"Keras prediction failed: {e}")
-                step2_out = step1_out
-        else:
-            logging.info("Keras model skipped due to missing or failed load")
+        subprocess.run(["iqtree2", "-s", aligned_file, "-nt", "AUTO"], check=True)
+    except Exception as e:
+        logging.error(f"IQ-TREE failed: {e}")
 
-        # Step 3: Run ML Simplified Tree using PhylogeneticTreeAnalyzer with existing CSV
+    # Step 3: ML Simplified Tree
+    html_file = "phylogenetic_tree_normalized_horizontal.html"
+    ml_output = ""
+    if os.path.exists(csv_path):
         analyzer = ml_simplified_tree.PhylogeneticTreeAnalyzer()
-        csv_path = "f gene clean dataset.csv"
-        if not os.path.exists(csv_path):
-            ml_output = f"Error: Expected CSV file '{csv_path}' not found"
-            logging.error(ml_output)
+        if analyzer.load_data(csv_path):
+            if analyzer.find_query_sequence(step2_out):
+                matched_ids, perc = analyzer.find_similar_sequences(analyzer.matching_percentage)
+                analyzer.create_interactive_tree(matched_ids, perc)
+                ml_output = "Tree generated."
+            else:
+                ml_output = "Query sequence not found."
         else:
-            try:
-                if not analyzer.load_data(csv_path):
-                    ml_output = "Error: Failed to load data from existing CSV"
-                    logging.error(f"Load data failed: Check CSV content. Head: {pd.read_csv(csv_path).head()}")
-                else:
-                    if not analyzer.find_query_sequence(step2_out):
-                        ml_output = "Error: Invalid query sequence"
-                        logging.error(f"Query sequence invalid: {step2_out[:50]}... (truncated)")
-                    else:
-                        matched_ids, actual_percentage = analyzer.find_similar_sequences(analyzer.matching_percentage)
-                        if not matched_ids:
-                            ml_output = "Error: No similar sequences found"
-                            logging.error(f"No matches for sequence: {step2_out[:50]}... (truncated)")
-                        else:
-                            analyzer.create_interactive_tree(matched_ids, actual_percentage)
-                            ml_output = f"Tree generated and saved as phylogenetic_tree_normalized_horizontal.html"
-                            if not os.path.exists("phylogenetic_tree_normalized_horizontal.html"):
-                                ml_output += " (File not found after generation)"
-                                logging.error("HTML file not generated")
-                            else:
-                                logging.info("HTML file generated successfully")
-            except Exception as e:
-                ml_output = f"Error in ML tree generation: {e}"
-                logging.error(f"ML tree generation failed: {e}")
-
-        html_file = "phylogenetic_tree_normalized_horizontal.html" if os.path.exists("phylogenetic_tree_normalized_horizontal.html") else None
-
-        return {
-            "Boundary Model Output": step1_out if boundary_model else "Skipped (Model not loaded)",
-            "Keras Model Output": step2_out if keras_model and kmer_to_index else "Skipped (Model not loaded)",
-            "CSV Save Status": "Using existing 'f gene clean.csv'",
-            "Maximum Likelihood Tree Output": ml_output,
-            "Tree HTML File": html_file
-        }
-    except Exception as e:
-        logging.error(f"Pipeline failed: {e}")
-        return {
-            "Boundary Model Output": f"Error: {e}",
-            "Keras Model Output": "N/A",
-            "CSV Save Status": "N/A",
-            "Maximum Likelihood Tree Output": "N/A",
-            "Tree HTML File": None
-        }
-
-# --------- Gradio Interface and API ---------
-with gr.Blocks() as gr_interface:
-    gr.Markdown("# Sequential Phylogenetic Inference Pipeline")
-    gr.Markdown("This pipeline generates a phylogenetic tree using the existing 'f gene clean.csv'. Model files are expected in the 'models/' subdirectory and are optional.")
-    
-    dna_input = gr.Textbox(label="Input DNA Sequence", placeholder="Enter DNA sequence (e.g., ATGCGTAACTAGCTAGCTAA)")
-    submit_button = gr.Button("Generate Tree")
-    
-    boundary_output = gr.Textbox(label="Boundary Model Output")
-    keras_output = gr.Textbox(label="Keras Model Output")
-    csv_status = gr.Textbox(label="CSV Save Status")
-    ml_output = gr.Textbox(label="Maximum Likelihood Tree Output")
-    tree_download = gr.File(label="Download Tree (HTML)", file_types=[".html"])
-    
-    submit_button.click(
-        fn=run_full_pipeline,
-        inputs=dna_input,
-        outputs=[
-            boundary_output,
-            keras_output,
-            csv_status,
-            ml_output,
-            tree_download
-        ]
-    )
-
-# --------- Launch ---------
-if __name__ == "__main__":
-    try:
-        gr_interface.launch(server_name="0.0.0.0", server_port=7860)
-        logging.info("Gradio interface launched successfully")
-    except Exception as e:
-        logging.error(f"Error launching Gradio interface: {e}")
-        raise
+            ml_output = "Failed to load CSV."
+    else:
+        ml_output = "CSV file missing."
+
+    return step1_out, step2_out, csv_path, ml_output, html_file, aligned_file, phy_file
+
+# --- Gradio UI ---
+with gr.Blocks() as demo:
+    gr.Markdown("# Viral Gene Phylogenetic Inference Pipeline")
+
+    inp = gr.Textbox(label="DNA Input")
+    btn = gr.Button("Run Pipeline")
+
+    out1 = gr.Textbox(label="Boundary Model Output")
+    out2 = gr.Textbox(label="Keras Model Output")
+    out3 = gr.Textbox(label="CSV File Used")
+    out4 = gr.Textbox(label="ML Tree Output")
+    html = gr.File(label="ML Tree (HTML)", file_types=['.html'])
+    fasta = gr.File(label="Aligned FASTA", file_types=['.fasta'])
+    phy = gr.File(label="IQ-TREE .phy File", file_types=['.phy'])
+
+    btn.click(fn=run_pipeline, inputs=inp, outputs=[out1, out2, out3, out4, html, fasta, phy])
+
+if __name__ == '__main__':
+    demo.launch(server_name="0.0.0.0", server_port=7860)