Spaces:
Sleeping
Sleeping
| """ | |
| Bioinformatics Protein Analysis Application | |
| Optimized for Hugging Face Spaces & Docker Deployment | |
| """ | |
| import os | |
| import io | |
| import csv | |
| import sys | |
| import subprocess | |
| import pandas as pd | |
| import joblib | |
| import numpy as np | |
| import torch | |
| import esm | |
| import logging | |
| from pathlib import Path | |
| from datetime import datetime | |
| from flask import Flask, request, render_template, send_file, session, redirect, url_for | |
| from Bio.Seq import Seq | |
| # ========== LOGGING SETUP ========== | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s %(levelname)s %(message)s", | |
| handlers=[ | |
| logging.FileHandler("app.log"), | |
| logging.StreamHandler(sys.stdout) | |
| ], | |
| ) | |
| logger = logging.getLogger(__name__) | |
| # ========== ENVIRONMENT & PATHS ========== | |
| # Get base directory (works in Docker and local) | |
| BASE_DIR = Path(os.getenv("BASE_DIR", "/app")) | |
| DATA_DIR = BASE_DIR / "data" | |
| TEMP_DIR = BASE_DIR / "temp" | |
| MODEL_DIR = BASE_DIR / "models" | |
| TEMPLATE_DIR = BASE_DIR / "templates" | |
| STATIC_DIR = BASE_DIR / "static" | |
| # Create directories if they don't exist | |
| for dir_path in [DATA_DIR, TEMP_DIR, MODEL_DIR]: | |
| dir_path.mkdir(parents=True, exist_ok=True) | |
| # Model paths - use defaults if not found | |
| MODEL_PATH = MODEL_DIR / "best_model_LR.sav" | |
| ENCODER_PATH = MODEL_DIR / "label_encoder.pkl" | |
| # BLAST configuration | |
| BLAST_PATH = os.getenv("BLAST_PATH", "/usr/bin/blastp") # Linux default | |
| BLAST_DB = DATA_DIR / "pathway_db" | |
| MAPPING_FILE = DATA_DIR / "pathway_map.csv" | |
| # Thresholds | |
| BITSCORE_THRESHOLD = 80.0 | |
| CONFIDENCE_THRESHOLD = 0.7 | |
| RESULT_CSV = TEMP_DIR / "combined_results.csv" | |
| # ========== MODEL LOADING ========== | |
| model = None | |
| encoder = None | |
| esm_model = None | |
| batch_converter = None | |
| logger.info("===== Application Startup =====") | |
| logger.info(f"BASE_DIR: {BASE_DIR}") | |
| logger.info(f"DATA_DIR: {DATA_DIR}") | |
| logger.info(f"TEMP_DIR: {TEMP_DIR}") | |
| # Load Logistic Regression Model | |
| try: | |
| if MODEL_PATH.exists(): | |
| model = joblib.load(MODEL_PATH) | |
| logger.info("✅ LR model loaded successfully.") | |
| else: | |
| logger.warning(f"⚠️ Model not found at {MODEL_PATH}") | |
| except Exception as e: | |
| logger.error(f"❌ Failed to load LR model: {e}") | |
| # Load Label Encoder | |
| try: | |
| if ENCODER_PATH.exists(): | |
| encoder = joblib.load(ENCODER_PATH) | |
| logger.info("✅ Label encoder loaded successfully.") | |
| else: | |
| logger.warning(f"⚠️ Encoder not found at {ENCODER_PATH}") | |
| except Exception as e: | |
| logger.error(f"❌ Failed to load label encoder: {e}") | |
| # Load ESM2 Model (downloads if not cached) | |
| try: | |
| logger.info("Loading ESM2 model (this may take a moment on first run)...") | |
| esm_model, alphabet = esm.pretrained.load_model_and_alphabet("esm2_t6_8M_UR50D") | |
| esm_model.eval() | |
| esm_model = esm_model.to("cpu") | |
| batch_converter = alphabet.get_batch_converter() | |
| logger.info("✅ ESM2 model loaded successfully on CPU.") | |
| except Exception as e: | |
| logger.error(f"❌ Failed to load ESM2 model: {e}") | |
| # ========== HELPER FUNCTIONS ========== | |
| def get_best_frame_protein(nuc_seq): | |
| """Translate DNA to best reading frame protein""" | |
| nuc_seq = "".join(nuc_seq.split()).upper() | |
| trim_len = len(nuc_seq) - (len(nuc_seq) % 3) | |
| nuc_seq = nuc_seq[:trim_len] | |
| seq = Seq(nuc_seq) | |
| frames = [ | |
| seq[i:].translate(to_stop=False) for i in range(3) | |
| ] + [ | |
| seq.reverse_complement()[i:].translate(to_stop=False) for i in range(3) | |
| ] | |
| best_protein = "" | |
| for prot in frames: | |
| longest_orf = max(prot.split("*"), key=len) | |
| if len(longest_orf) > len(best_protein): | |
| best_protein = longest_orf | |
| return str(best_protein) | |
| def detect_and_translate(seq): | |
| """Detect if DNA or protein and handle accordingly""" | |
| seq = seq.upper() | |
| # Check if DNA | |
| if all(base in "ATGCN" for base in seq): | |
| logger.info("Detected DNA sequence → translating to protein.") | |
| if len(seq) < 150: | |
| logger.warning("DNA sequence too short for translation") | |
| return None | |
| prot_seq = get_best_frame_protein(seq) | |
| if len(prot_seq) < 50: | |
| logger.warning("Translated protein too short") | |
| return None | |
| return prot_seq | |
| else: | |
| logger.info("Detected protein sequence → using directly.") | |
| if len(seq) < 50: | |
| logger.warning("Protein sequence too short") | |
| return None | |
| return seq | |
| def esm2_320_embed(sequence): | |
| """Generate ESM2 embeddings for sequence""" | |
| if esm_model is None or batch_converter is None: | |
| raise RuntimeError("ESM model not available.") | |
| batch_labels, batch_strs, batch_tokens = batch_converter([("seq1", sequence)]) | |
| with torch.no_grad(): | |
| results = esm_model(batch_tokens, repr_layers=[6], return_contacts=False) | |
| token_representations = results["representations"][6] | |
| seq_repr = token_representations[0, 1:-1].detach().cpu().numpy() | |
| return seq_repr.mean(axis=0) | |
| def run_blast_and_get_dataframe(temp_fasta, blast_output): | |
| """Execute BLAST and return results as DataFrame""" | |
| try: | |
| # Check if BLAST is available | |
| if not Path(BLAST_PATH).exists(): | |
| logger.warning(f"⚠️ BLAST not found at {BLAST_PATH}, skipping BLAST search") | |
| return pd.DataFrame( | |
| columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"] | |
| ) | |
| # Check if database exists | |
| if not BLAST_DB.exists(): | |
| logger.warning(f"⚠️ BLAST database not found at {BLAST_DB}, skipping BLAST search") | |
| return pd.DataFrame( | |
| columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"] | |
| ) | |
| cmd = [ | |
| str(BLAST_PATH), | |
| "-query", str(temp_fasta), | |
| "-db", str(BLAST_DB), | |
| "-out", str(blast_output), | |
| "-outfmt", "6 qseqid sseqid pident length evalue bitscore" | |
| ] | |
| logger.info(f"🧬 Running BLAST...") | |
| proc = subprocess.run(cmd, capture_output=True, text=True, timeout=300) | |
| if proc.returncode != 0: | |
| logger.error(f"❌ BLAST error: {proc.stderr}") | |
| return pd.DataFrame( | |
| columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"] | |
| ) | |
| if not Path(blast_output).exists() or Path(blast_output).stat().st_size == 0: | |
| logger.warning("⚠️ No BLAST results generated.") | |
| return pd.DataFrame( | |
| columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"] | |
| ) | |
| cols = ["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"] | |
| df = pd.read_csv(blast_output, sep="\t", names=cols) | |
| df["bitscore"] = df["bitscore"].astype(float) | |
| logger.info(f"✅ BLAST completed successfully with {len(df)} hits") | |
| return df | |
| except subprocess.TimeoutExpired: | |
| logger.error("❌ BLAST timed out (>300s)") | |
| return pd.DataFrame( | |
| columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"] | |
| ) | |
| except Exception as e: | |
| logger.error(f"❌ Unexpected error in BLAST: {e}") | |
| return pd.DataFrame( | |
| columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"] | |
| ) | |
| # ========== FLASK APPLICATION ========== | |
| app = Flask( | |
| __name__, | |
| template_folder=str(TEMPLATE_DIR), | |
| static_folder=str(STATIC_DIR) | |
| ) | |
| app.secret_key = os.getenv("SECRET_KEY", "default-secret-key-change-in-production") | |
| app.config['ENV'] = os.getenv("FLASK_ENV", "production") | |
| app.config['DEBUG'] = os.getenv("FLASK_DEBUG", "False") == "True" | |
| latest_predictions = [] | |
| def add_cache_control(response): | |
| """Prevent caching of responses""" | |
| response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate" | |
| response.headers["Pragma"] = "no-cache" | |
| response.headers["Expires"] = "0" | |
| return response | |
| def index(): | |
| """Main index page - handle sequence input and analysis""" | |
| global latest_predictions | |
| latest_predictions = [] | |
| if request.method == "POST": | |
| sequences = [] | |
| headers = [] | |
| # Parse uploaded FASTA file | |
| uploaded_file = request.files.get("fasta_file") | |
| if uploaded_file and uploaded_file.filename != "": | |
| seq = "" | |
| header = "" | |
| for line in uploaded_file: | |
| line = line.decode().strip() | |
| if not line: | |
| continue | |
| if line.startswith(">"): | |
| if seq: | |
| sequences.append(seq) | |
| headers.append(header) | |
| seq = "" | |
| header = line[1:] | |
| else: | |
| seq += line | |
| if seq: | |
| sequences.append(seq) | |
| headers.append(header) | |
| # Parse textarea input | |
| sequence_text = request.form.get("sequence_text", "").strip() | |
| if sequence_text: | |
| seq = "" | |
| header = "" | |
| for line in sequence_text.splitlines(): | |
| line = line.strip() | |
| if not line: | |
| continue | |
| if line.startswith(">"): | |
| if seq: | |
| sequences.append(seq) | |
| headers.append(header) | |
| seq = "" | |
| header = line[1:] | |
| else: | |
| seq += line | |
| if seq: | |
| sequences.append(seq) | |
| headers.append(header or "sequence_from_text") | |
| if not sequences: | |
| return render_template( | |
| "index.html", | |
| error="⚠️ No valid sequences provided.", | |
| current_year=datetime.now().year | |
| ) | |
| # Detect and translate sequences | |
| translated_sequences = [detect_and_translate(s) for s in sequences] | |
| filtered_sequences = [s for s in translated_sequences if s is not None] | |
| filtered_headers = [ | |
| h for s, h in zip(translated_sequences, headers) if s is not None | |
| ] | |
| if not filtered_sequences: | |
| return render_template( | |
| "index.html", | |
| error="⚠️ All sequences too short or invalid.", | |
| current_year=datetime.now().year | |
| ) | |
| # Gene Family Prediction using ESM2 embeddings | |
| gene_preds = [] | |
| try: | |
| logger.info(f"Generating embeddings for {len(filtered_sequences)} sequences...") | |
| features = np.array([esm2_320_embed(seq) for seq in filtered_sequences]) | |
| if model and encoder: | |
| probs = model.predict_proba(features) | |
| max_probs = probs.max(axis=1) | |
| pred_indices = probs.argmax(axis=1) | |
| for idx, prob in zip(pred_indices, max_probs): | |
| if prob >= CONFIDENCE_THRESHOLD: | |
| gene_preds.append(encoder.inverse_transform([idx])[0]) | |
| else: | |
| gene_preds.append("Unknown") | |
| else: | |
| logger.warning("Model or encoder not available, marking all as Unknown") | |
| gene_preds = ["Unknown"] * len(filtered_sequences) | |
| except Exception as e: | |
| logger.error(f"Gene family prediction failed: {e}") | |
| gene_preds = ["Unknown"] * len(filtered_sequences) | |
| # Save temporary FASTA for BLAST | |
| temp_fasta = TEMP_DIR / "temp_input.fasta" | |
| with open(temp_fasta, "w") as f: | |
| for h, s in zip(filtered_headers, filtered_sequences): | |
| f.write(f">{h.replace(' ', '_')}\n{s}\n") | |
| # Run BLAST | |
| blast_output = TEMP_DIR / "blast_results.txt" | |
| blast_df = run_blast_and_get_dataframe(temp_fasta, blast_output) | |
| # Load pathway mapping file | |
| map_df = pd.DataFrame() | |
| if MAPPING_FILE.exists(): | |
| try: | |
| map_df = pd.read_csv(MAPPING_FILE, dtype=str) | |
| logger.info(f"Loaded {len(map_df)} pathway mappings") | |
| except Exception as e: | |
| logger.error(f"Failed to read mapping file: {e}") | |
| # Build results | |
| combined_results = [] | |
| if not blast_df.empty: | |
| top_hits_idx = blast_df.groupby("qseqid")["bitscore"].idxmax() | |
| top_hits = blast_df.loc[top_hits_idx] | |
| else: | |
| top_hits = pd.DataFrame( | |
| columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"] | |
| ) | |
| for i, orig_header in enumerate(filtered_headers): | |
| safe_h = orig_header.replace(" ", "_") | |
| top_hit_row = top_hits[top_hits["qseqid"] == safe_h] | |
| if not top_hit_row.empty: | |
| row = top_hit_row.iloc[0] | |
| top_hit = row["sseqid"] | |
| bitscore = float(row["bitscore"]) | |
| pathways = "No Pathways Found" | |
| if ( | |
| not map_df.empty | |
| and "Entry" in map_df.columns | |
| and "Pathways" in map_df.columns | |
| ): | |
| match = map_df.loc[map_df["Entry"] == top_hit, "Pathways"] | |
| if not match.empty and bitscore >= BITSCORE_THRESHOLD: | |
| pathways = match.values[0] | |
| else: | |
| top_hit = "No Hit" | |
| bitscore = 0.0 | |
| pathways = "No Pathways Found" | |
| gene_family = gene_preds[i] if i < len(gene_preds) else "Unknown" | |
| combined_results.append([orig_header, top_hit, bitscore, pathways, gene_family]) | |
| # Save results to CSV | |
| result_df = pd.DataFrame( | |
| combined_results, | |
| columns=["Query/Header", "Top Hit", "Bitscore", "Pathways", "Predicted Gene Family"] | |
| ) | |
| result_df.to_csv(RESULT_CSV, index=False) | |
| latest_predictions = combined_results | |
| # Store in session | |
| session['predictions'] = combined_results | |
| session['result_summary'] = { | |
| 'total_sequences': len(filtered_sequences), | |
| 'successful_predictions': len([p for p in gene_preds if p != "Unknown"]), | |
| 'total_pathways_found': len([r for r in combined_results if r[3] != "No Pathways Found"]) | |
| } | |
| return redirect(url_for('results')) | |
| return render_template("index.html", predictions=[], current_year=datetime.now().year) | |
| def results(): | |
| """Display analysis results""" | |
| predictions = session.get("predictions", []) | |
| summary = session.get("result_summary", {}) | |
| if not predictions: | |
| return redirect(url_for("index")) | |
| return render_template( | |
| "results.html", | |
| predictions=predictions, | |
| summary=summary, | |
| current_year=datetime.now().year | |
| ) | |
| def download_csv(): | |
| """Download results as CSV""" | |
| global latest_predictions | |
| if not latest_predictions: | |
| return "⚠️ No predictions to download." | |
| output = io.StringIO() | |
| writer = csv.writer(output) | |
| writer.writerow(["Query/Header", "Top Hit", "Bitscore", "Pathways", "Predicted Gene Family"]) | |
| writer.writerows(latest_predictions) | |
| output.seek(0) | |
| return send_file( | |
| io.BytesIO(output.getvalue().encode()), | |
| mimetype="text/csv", | |
| as_attachment=True, | |
| download_name="bioinformatics_results.csv" | |
| ) | |
| def health(): | |
| """Health check endpoint""" | |
| return { | |
| "status": "healthy", | |
| "esm_model": esm_model is not None, | |
| "lr_model": model is not None, | |
| "encoder": encoder is not None | |
| } | |
| # ========== ERROR HANDLERS ========== | |
| def not_found(error): | |
| return render_template("error.html", error="Page not found"), 404 | |
| def internal_error(error): | |
| logger.error(f"Internal error: {error}") | |
| return render_template("error.html", error="Internal server error"), 500 | |
| # ========== MAIN ========== | |
| if __name__ == "__main__": | |
| # Get port from environment or default to 7860 (HF Spaces) | |
| port = int(os.getenv("PORT", 7860)) | |
| host = os.getenv("HOST", "0.0.0.0") | |
| debug = os.getenv("FLASK_DEBUG", "False") == "True" | |
| logger.info(f"Starting Flask app on {host}:{port} (debug={debug})") | |
| app.run( | |
| host=host, | |
| port=port, | |
| debug=debug, | |
| threaded=True, | |
| use_reloader=False | |
| ) |