Update app.py

app.py

@@ -6,49 +6,109 @@ import logging
 import traceback
 from pathlib import Path
 
-# Import your pipeline modules
-try:
-    from utils_audio import convert_to_wav
-    from to_cha import to_cha_from_wav
-    from cha_json import cha_to_json_file
-    from output import predict_from_chajson
-except ImportError as e:
-    logging.error(f"Import error: {e}")
-    # Fallback imports or error handling
-
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Configuration
-MODEL_DIR = "./adaptive_aphasia_model"  # Path to your trained model
+# Configuration - Use current directory for model files
+MODEL_DIR = "."  # Changed from "./adaptive_aphasia_model"
 SUPPORTED_AUDIO_FORMATS = [".mp3", ".mp4", ".wav", ".m4a", ".flac", ".ogg"]
 
+def safe_import_modules():
+    """Safely import pipeline modules with error handling"""
+    modules = {}
+    
+    try:
+        from utils_audio import convert_to_wav
+        modules['convert_to_wav'] = convert_to_wav
+        logger.info("✓ utils_audio imported successfully")
+    except Exception as e:
+        logger.error(f"✗ Failed to import utils_audio: {e}")
+        modules['convert_to_wav'] = None
+    
+    try:
+        from to_cha import to_cha_from_wav
+        modules['to_cha_from_wav'] = to_cha_from_wav
+        logger.info("✓ to_cha imported successfully")
+    except Exception as e:
+        logger.error(f"✗ Failed to import to_cha: {e}")
+        modules['to_cha_from_wav'] = None
+    
+    try:
+        from cha_json import cha_to_json_file
+        modules['cha_to_json_file'] = cha_to_json_file
+        logger.info("✓ cha_json imported successfully")
+    except Exception as e:
+        logger.error(f"✗ Failed to import cha_json: {e}")
+        modules['cha_to_json_file'] = None
+    
+    try:
+        from output import predict_from_chajson
+        modules['predict_from_chajson'] = predict_from_chajson
+        logger.info("✓ output imported successfully")
+    except Exception as e:
+        logger.error(f"✗ Failed to import output: {e}")
+        modules['predict_from_chajson'] = None
+    
+    return modules
+
+# Import modules
+MODULES = safe_import_modules()
+
+def check_model_files():
+    """Check if required model files exist"""
+    required_files = [
+        "pytorch_model.bin",
+        "config.json",
+        "tokenizer.json",
+        "tokenizer_config.json"
+    ]
+    
+    missing_files = []
+    for file in required_files:
+        if not os.path.exists(os.path.join(MODEL_DIR, file)):
+            missing_files.append(file)
+    
+    if missing_files:
+        logger.error(f"Missing model files: {missing_files}")
+        return False, missing_files
+    
+    logger.info("✓ All required model files found")
+    return True, []
+
 def run_complete_pipeline(audio_file_path: str) -> dict:
-    """
-    Complete pipeline: Audio → WAV → CHA → JSON → Model Prediction
-    """
+    """Complete pipeline: Audio → WAV → CHA → JSON → Model Prediction"""
+    
+    # Check if all modules are available
+    if not all(MODULES.values()):
+        missing = [k for k, v in MODULES.items() if v is None]
+        return {
+            "success": False,
+            "error": f"Missing required modules: {missing}",
+            "message": "Pipeline modules not available"
+        }
+    
     try:
         logger.info(f"Starting pipeline for: {audio_file_path}")
         
         # Step 1: Convert to WAV
         logger.info("Step 1: Converting audio to WAV...")
-        wav_path = convert_to_wav(audio_file_path, sr=16000, mono=True)
+        wav_path = MODULES['convert_to_wav'](audio_file_path, sr=16000, mono=True)
         logger.info(f"WAV conversion completed: {wav_path}")
         
         # Step 2: Generate CHA file using Batchalign
         logger.info("Step 2: Generating CHA file...")
-        cha_path = to_cha_from_wav(wav_path, lang="eng")
+        cha_path = MODULES['to_cha_from_wav'](wav_path, lang="eng")
         logger.info(f"CHA generation completed: {cha_path}")
         
         # Step 3: Convert CHA to JSON
         logger.info("Step 3: Converting CHA to JSON...")
-        chajson_path, json_data = cha_to_json_file(cha_path)
+        chajson_path, json_data = MODULES['cha_to_json_file'](cha_path)
         logger.info(f"JSON conversion completed: {chajson_path}")
         
         # Step 4: Run aphasia classification
         logger.info("Step 4: Running aphasia classification...")
-        results = predict_from_chajson(MODEL_DIR, chajson_path, output_file=None)
+        results = MODULES['predict_from_chajson'](MODEL_DIR, chajson_path, output_file=None)
         logger.info("Classification completed")
         
         # Cleanup temporary files
@@ -75,82 +135,44 @@ def run_complete_pipeline(audio_file_path: str) -> dict:
         }
 
 def process_audio_input(audio_file):
-    """
-    Process audio file and return formatted results
-    """
+    """Process audio file and return formatted results"""
     try:
         if audio_file is None:
-            return (
-                "❌ Error: No audio file uploaded",
-                "",
-                "",
-                "",
-                ""
-            )
+            return "❌ Error: No audio file uploaded"
+        
+        # Check if pipeline is available
+        if not all(MODULES.values()):
+            return "❌ Error: Audio processing pipeline not available. Missing required modules."
         
         # Check file format
         file_path = audio_file
-        if isinstance(audio_file, str):
-            file_path = audio_file
-        else:
-            # Handle Gradio file object
-            file_path = audio_file.name if hasattr(audio_file, 'name') else str(audio_file)
+        if hasattr(audio_file, 'name'):
+            file_path = audio_file.name
         
         file_ext = Path(file_path).suffix.lower()
         if file_ext not in SUPPORTED_AUDIO_FORMATS:
-            return (
-                f"❌ Error: Unsupported file format {file_ext}",
-                f"Supported formats: {', '.join(SUPPORTED_AUDIO_FORMATS)}",
-                "",
-                "",
-                ""
-            )
+            return f"❌ Error: Unsupported file format {file_ext}. Supported: {', '.join(SUPPORTED_AUDIO_FORMATS)}"
         
         # Run the complete pipeline
         pipeline_result = run_complete_pipeline(file_path)
         
         if not pipeline_result["success"]:
-            return (
-                f"❌ Pipeline Error: {pipeline_result['message']}",
-                pipeline_result.get('error', ''),
-                "",
-                "",
-                ""
-            )
+            return f"❌ Pipeline Error: {pipeline_result['message']}\n\nDetails: {pipeline_result.get('error', '')}"
         
-        # Extract results
+        # Format results
         results = pipeline_result["results"]
         
-        # Format main prediction
         if "predictions" in results and len(results["predictions"]) > 0:
             first_pred = results["predictions"][0]
             
             if "error" in first_pred:
-                return (
-                    f"❌ Classification Error: {first_pred['error']}",
-                    "",
-                    "",
-                    "",
-                    ""
-                )
+                return f"❌ Classification Error: {first_pred['error']}"
             
-            # Main prediction
+            # Format main result
             predicted_class = first_pred["prediction"]["predicted_class"]
             confidence = first_pred["prediction"]["confidence_percentage"]
-            class_description = first_pred["class_description"]["name"]
-            
-            main_result = f"🧠 **Predicted Aphasia Type:** {predicted_class}\n"
-            main_result += f"📊 **Confidence:** {confidence}\n"
-            main_result += f"📋 **Description:** {class_description}"
-            
-            # Detailed analysis
-            features = first_pred["class_description"].get("features", [])
-            detailed_analysis = f"**Key Features:**\n"
-            for feature in features:
-                detailed_analysis += f"• {feature}\n"
-            
-            detailed_analysis += f"\n**Clinical Description:**\n"
-            detailed_analysis += first_pred["class_description"].get("description", "No description available")
+            class_name = first_pred["class_description"]["name"]
+            description = first_pred["class_description"]["description"]
             
             # Additional metrics
             additional_info = first_pred["additional_predictions"]
@@ -158,65 +180,55 @@ def process_audio_input(audio_file):
             fluency_score = additional_info["fluency_score"]
             fluency_rating = additional_info["fluency_rating"]
             
-            additional_metrics = f"**Severity Level:** {severity_level}/3\n"
-            additional_metrics += f"**Fluency Score:** {fluency_score:.3f} ({fluency_rating})\n"
-            
-            # Probability distribution (top 3)
+            # Format probability distribution (top 3)
             prob_dist = first_pred["probability_distribution"]
             top_3 = list(prob_dist.items())[:3]
             
-            probability_breakdown = "**Top 3 Classifications:**\n"
+            result_text = f"""
+🧠 **APHASIA CLASSIFICATION RESULTS**
+
+🎯 **Primary Classification:** {predicted_class}
+📊 **Confidence:** {confidence}
+📋 **Type:** {class_name}
+
+📈 **Additional Metrics:**
+• Severity Level: {severity_level}/3
+• Fluency Score: {fluency_score:.3f} ({fluency_rating})
+
+📊 **Top 3 Probability Rankings:**
+"""
             for i, (aphasia_type, info) in enumerate(top_3, 1):
-                probability_breakdown += f"{i}. {aphasia_type}: {info['percentage']}\n"
+                result_text += f"{i}. {aphasia_type}: {info['percentage']}\n"
             
-            # Summary statistics
-            summary = results.get("summary", {})
-            summary_text = f"**Processing Summary:**\n"
-            summary_text += f"• Total sentences analyzed: {results.get('total_sentences', 'N/A')}\n"
-            summary_text += f"• Average confidence: {summary.get('average_confidence', 'N/A')}\n"
-            summary_text += f"• Average fluency: {summary.get('average_fluency_score', 'N/A')}\n"
+            result_text += f"""
+📝 **Clinical Description:**
+{description}
+
+📊 **Processing Summary:**
+• Total sentences analyzed: {results.get('total_sentences', 'N/A')}
+• Average confidence: {results.get('summary', {}).get('average_confidence', 'N/A')}
+• Average fluency: {results.get('summary', {}).get('average_fluency_score', 'N/A')}
+"""
             
-            return (
-                main_result,
-                detailed_analysis,
-                additional_metrics,
-                probability_breakdown,
-                summary_text
-            )
+            return result_text
         
         else:
-            return (
-                "❌ No predictions generated",
-                "The audio file may not contain analyzable speech",
-                "",
-                "",
-                ""
-            )
+            return "❌ No predictions generated. The audio file may not contain analyzable speech."
             
     except Exception as e:
         logger.error(f"Processing error: {str(e)}")
         logger.error(traceback.format_exc())
-        return (
-            f"❌ Processing Error: {str(e)}",
-            "Please check the logs for more details",
-            "",
-            "",
-            ""
-        )
+        return f"❌ Processing Error: {str(e)}\n\nPlease check the logs for more details."
 
 def process_text_input(text_input):
-    """
-    Process text input directly (fallback option)
-    """
+    """Process text input directly (fallback option)"""
     try:
         if not text_input or not text_input.strip():
-            return (
-                "❌ Error: Please enter some text for analysis",
-                "",
-                "",
-                "",
-                ""
-            )
+            return "❌ Error: Please enter some text for analysis"
+        
+        # Check if prediction module is available
+        if MODULES['predict_from_chajson'] is None:
+            return "❌ Error: Text analysis not available. Missing prediction module."
         
         # Create a simple JSON structure for text-only input
         temp_json = {
@@ -243,7 +255,7 @@ def process_text_input(text_input):
             temp_json_path = f.name
         
         # Run prediction
-        results = predict_from_chajson(MODEL_DIR, temp_json_path, output_file=None)
+        results = MODULES['predict_from_chajson'](MODEL_DIR, temp_json_path, output_file=None)
         
         # Cleanup
         try:
@@ -251,227 +263,134 @@ def process_text_input(text_input):
         except:
             pass
         
-        # Format results (similar to audio processing)
+        # Format results
         if "predictions" in results and len(results["predictions"]) > 0:
             first_pred = results["predictions"][0]
             
             predicted_class = first_pred["prediction"]["predicted_class"]
             confidence = first_pred["prediction"]["confidence_percentage"]
+            description = first_pred["class_description"]["description"]
+            severity = first_pred["additional_predictions"]["predicted_severity_level"]
+            fluency = first_pred["additional_predictions"]["fluency_rating"]
             
-            return (
-                f"🧠 **Predicted:** {predicted_class} ({confidence})",
-                first_pred["class_description"]["description"],
-                f"Severity: {first_pred['additional_predictions']['predicted_severity_level']}/3",
-                f"Fluency: {first_pred['additional_predictions']['fluency_rating']}",
-                "Text-based analysis completed"
-            )
+            return f"""
+🧠 **TEXT ANALYSIS RESULTS**
+
+🎯 **Predicted:** {predicted_class}
+📊 **Confidence:** {confidence}
+📈 **Severity:** {severity}/3
+🗣️ **Fluency:** {fluency}
+
+📝 **Description:**
+{description}
+
+ℹ️ **Note:** Text-based analysis provides limited accuracy compared to audio analysis.
+"""
         else:
-            return (
-                "❌ No predictions generated",
-                "",
-                "",
-                "",
-                ""
-            )
+            return "❌ No predictions generated from text input"
     
     except Exception as e:
         logger.error(f"Text processing error: {str(e)}")
-        return (
-            f"❌ Error: {str(e)}",
-            "",
-            "",
-            "",
-            ""
-        )
+        return f"❌ Error: {str(e)}"
 
-# Create Gradio interface
 def create_interface():
-    """Create the main Gradio interface"""
+    """Create simplified Gradio interface"""
+    
+    # Check system status
+    model_available, missing_files = check_model_files()
+    pipeline_available = all(MODULES.values())
     
-    with gr.Blocks(
-        title="Advanced Aphasia Classification System",
-        theme=gr.themes.Soft(),
-        css="""
-        .main-header { text-align: center; margin-bottom: 2rem; }
-        .upload-section { border: 2px dashed #ccc; padding: 2rem; border-radius: 10px; }
-        .results-section { margin-top: 2rem; }
-        """
-    ) as demo:
+    status_message = "🟢 **System Status: Ready**" if model_available and pipeline_available else "🔴 **System Status: Issues Detected**"
+    
+    if not model_available:
+        status_message += f"\n❌ Missing model files: {', '.join(missing_files)}"
+    
+    if not pipeline_available:
+        missing_modules = [k for k, v in MODULES.items() if v is None]
+        status_message += f"\n❌ Missing modules: {', '.join(missing_modules)}"
+    
+    # Create interface
+    with gr.Blocks(title="Aphasia Classification System") as demo:
         
-        # Header
         gr.HTML("""
-        <div class="main-header">
+        <div style="text-align: center; margin-bottom: 2rem;">
             <h1>🧠 Advanced Aphasia Classification System</h1>
-            <p>Upload audio files (MP3, MP4, WAV) or enter text to analyze speech patterns and classify aphasia types</p>
+            <p>Upload audio files or enter text to analyze speech patterns and classify aphasia types</p>
         </div>
         """)
         
+        gr.Markdown(status_message)
+        
         with gr.Tabs():
-            # Audio Input Tab
-            with gr.TabItem("🎵 Audio Analysis", id="audio_tab"):
+            # Audio Tab
+            with gr.Tab("🎵 Audio Analysis"):
                 gr.Markdown("### Upload Audio File")
-                gr.Markdown("Supported formats: MP3, MP4, WAV, M4A, FLAC, OGG")
+                gr.Markdown("**Supported formats:** MP3, MP4, WAV, M4A, FLAC, OGG")
                 
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        audio_input = gr.File(
-                            label="Upload Audio File",
-                            file_types=["audio"],
-                            type="filepath"
-                        )
-                        
-                        process_audio_btn = gr.Button(
-                            "🔍 Analyze Audio", 
-                            variant="primary",
-                            size="lg"
-                        )
-                        
-                        gr.Markdown("**Note:** Processing may take 1-3 minutes depending on audio length")
+                audio_input = gr.File(
+                    label="Upload Audio File",
+                    file_types=["audio"]
+                )
+                
+                audio_btn = gr.Button("🔍 Analyze Audio", variant="primary")
+                
+                audio_output = gr.Textbox(
+                    label="Analysis Results",
+                    lines=20,
+                    max_lines=30
+                )
                 
-                # Results section for audio
-                with gr.Column(scale=2, visible=True) as audio_results:
-                    gr.Markdown("### 📊 Analysis Results")
-                    
-                    audio_main_result = gr.Textbox(
-                        label="🎯 Primary Classification",
-                        lines=3,
-                        interactive=False
-                    )
-                    
-                    with gr.Row():
-                        audio_detailed = gr.Textbox(
-                            label="📋 Detailed Analysis",
-                            lines=6,
-                            interactive=False
-                        )
-                        
-                        audio_metrics = gr.Textbox(
-                            label="📈 Additional Metrics",
-                            lines=6,
-                            interactive=False
-                        )
-                    
-                    with gr.Row():
-                        audio_probabilities = gr.Textbox(
-                            label="📊 Probability Breakdown",
-                            lines=4,
-                            interactive=False
-                        )
-                        
-                        audio_summary = gr.Textbox(
-                            label="📝 Processing Summary",
-                            lines=4,
-                            interactive=False
-                        )
+                audio_btn.click(
+                    fn=process_audio_input,
+                    inputs=audio_input,
+                    outputs=audio_output
+                )
             
-            # Text Input Tab (Fallback)
-            with gr.TabItem("📝 Text Analysis", id="text_tab"):
+            # Text Tab
+            with gr.Tab("📝 Text Analysis"):
                 gr.Markdown("### Direct Text Input")
-                gr.Markdown("Enter speech transcription or text for analysis (fallback option)")
+                gr.Markdown("**Note:** Audio analysis provides more accurate results")
                 
-                with gr.Row():
-                    with gr.Column():
-                        text_input = gr.Textbox(
-                            label="Input Text",
-                            placeholder="Enter speech transcription or text for analysis...",
-                            lines=5
-                        )
-                        
-                        process_text_btn = gr.Button(
-                            "🔍 Analyze Text", 
-                            variant="secondary",
-                            size="lg"
-                        )
+                text_input = gr.Textbox(
+                    label="Enter Text",
+                    placeholder="Enter speech transcription or text for analysis...",
+                    lines=5
+                )
                 
-                # Results section for text
-                with gr.Column() as text_results:
-                    gr.Markdown("### 📊 Analysis Results")
-                    
-                    text_main_result = gr.Textbox(
-                        label="🎯 Primary Classification",
-                        lines=2,
-                        interactive=False
-                    )
-                    
-                    with gr.Row():
-                        text_detailed = gr.Textbox(
-                            label="📋 Clinical Description",
-                            lines=4,
-                            interactive=False
-                        )
-                        
-                        text_metrics = gr.Textbox(
-                            label="📈 Metrics",
-                            lines=4,
-                            interactive=False
-                        )
-                    
-                    with gr.Row():
-                        text_probabilities = gr.Textbox(
-                            label="📊 Assessment",
-                            lines=2,
-                            interactive=False
-                        )
-                        
-                        text_summary = gr.Textbox(
-                            label="📝 Status",
-                            lines=2,
-                            interactive=False
-                        )
-        
-        # Event handlers
-        process_audio_btn.click(
-            fn=process_audio_input,
-            inputs=[audio_input],
-            outputs=[
-                audio_main_result,
-                audio_detailed,
-                audio_metrics,
-                audio_probabilities,
-                audio_summary
-            ]
-        )
-        
-        process_text_btn.click(
-            fn=process_text_input,
-            inputs=[text_input],
-            outputs=[
-                text_main_result,
-                text_detailed,
-                text_metrics,
-                text_probabilities,
-                text_summary
-            ]
-        )
+                text_btn = gr.Button("🔍 Analyze Text", variant="secondary")
+                
+                text_output = gr.Textbox(
+                    label="Analysis Results",
+                    lines=15,
+                    max_lines=20
+                )
+                
+                text_btn.click(
+                    fn=process_text_input,
+                    inputs=text_input,
+                    outputs=text_output
+                )
         
-        # Footer
         gr.HTML("""
         <div style="text-align: center; margin-top: 2rem; padding: 1rem; border-top: 1px solid #eee;">
-            <p><strong>About:</strong> This system uses advanced NLP and acoustic analysis to classify different types of aphasia from speech samples.</p>
+            <p><strong>About:</strong> This system uses advanced NLP and acoustic analysis to classify different types of aphasia.</p>
             <p><em>For research and clinical assessment purposes.</em></p>
         </div>
         """)
     
     return demo
 
-# Launch the application
 if __name__ == "__main__":
     try:
         logger.info("Starting Aphasia Classification System...")
         
-        # Check if model directory exists
-        if not os.path.exists(MODEL_DIR):
-            logger.error(f"Model directory not found: {MODEL_DIR}")
-            print(f"❌ Error: Model directory not found: {MODEL_DIR}")
-            print("Please ensure your trained model is in the correct directory.")
-        
         # Create and launch interface
         demo = create_interface()
         demo.launch(
             server_name="0.0.0.0",
             server_port=7860,
-            show_error=True,
-            share=False
+            share=True,  # This fixes the localhost error
+            show_error=True
         )
         
     except Exception as e: