app.py

import torch, torchaudio, torchvision
import os
import gradio as gr
import numpy as np

from preprocess import process_audio_data, process_image_data
from train import WatermelonModel
from infer import infer

def load_model(model_path):
    global device
    device = torch.device(
        "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
    )
    print(f"\033[92mINFO\033[0m: Using device: {device}")

    # Check if the file exists
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file not found at {model_path}")
    
    # Check if the file is empty or very small
    file_size = os.path.getsize(model_path)
    if file_size < 1000:  # Less than 1KB is suspiciously small for a model
        print(f"\033[93mWARNING\033[0m: Model file size is only {file_size} bytes, which is suspiciously small")
    
    try:
        model = WatermelonModel().to(device)
        model.load_state_dict(torch.load(model_path, map_location=device))
        model.eval()
        print(f"\033[92mINFO\033[0m: Loaded model from {model_path}")
        return model
    except RuntimeError as e:
        if "failed finding central directory" in str(e):
            print(f"\033[91mERROR\033[0m: The model file at {model_path} appears to be corrupted.")
            print("This can happen if:")
            print("  1. The model saving process was interrupted")
            print("  2. The file was not properly downloaded")
            print("  3. The path points to a file that is not a valid PyTorch model")
            print(f"File size: {file_size} bytes")
        raise

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Watermelon sweetness predictor")
    parser.add_argument("--model_path", type=str, default="./models/model_15_20250405-033557.pt", help="Path to the trained model")
    args = parser.parse_args()

    model = load_model(args.model_path)

    def predict(audio, image):
        try:
            # Debug audio input
            print(f"\033[92mDEBUG\033[0m: Audio input type: {type(audio)}")
            print(f"\033[92mDEBUG\033[0m: Audio input value: {audio}")
            
            # Handle different formats of audio input from Gradio
            if audio is None:
                return "Error: No audio provided. Please upload or record audio."
                
            if isinstance(audio, tuple) and len(audio) >= 2:
                sr, audio_data = audio[0], audio[-1]
                print(f"\033[92mDEBUG\033[0m: Audio format: sr={sr}, audio_data shape={audio_data.shape if hasattr(audio_data, 'shape') else 'no shape'}")
            elif isinstance(audio, tuple) and len(audio) == 1:
                # Handle single element tuple
                audio_data = audio[0]
                sr = 44100  # Assume default sample rate
                print(f"\033[92mDEBUG\033[0m: Single element audio tuple, using default sr={sr}")
            elif isinstance(audio, np.ndarray):
                # Handle direct numpy array
                audio_data = audio
                sr = 44100  # Assume default sample rate
                print(f"\033[92mDEBUG\033[0m: Audio is numpy array, using default sr={sr}")
            else:
                return f"Error: Unexpected audio format: {type(audio)}"
            
            # Ensure audio_data is correctly shaped
            if isinstance(audio_data, np.ndarray):
                # Make sure we have a 2D array
                if len(audio_data.shape) == 1:
                    audio_data = np.expand_dims(audio_data, axis=0)
                    print(f"\033[92mDEBUG\033[0m: Reshaped 1D audio to 2D: {audio_data.shape}")
                
                # If channels are the second dimension, transpose
                if len(audio_data.shape) == 2 and audio_data.shape[0] > audio_data.shape[1]:
                    audio_data = np.transpose(audio_data)
                    print(f"\033[92mDEBUG\033[0m: Transposed audio shape to: {audio_data.shape}")
            
            # Convert to tensor
            audio_tensor = torch.tensor(audio_data).float()
            print(f"\033[92mDEBUG\033[0m: Audio tensor shape: {audio_tensor.shape}")
            
            # Process audio data and handle None case
            mfcc = process_audio_data(audio_tensor, sr)
            if mfcc is None:
                return "Error: Failed to process audio data. Make sure your audio contains a clear tapping sound."
                
            mfcc = mfcc.to(device)
            print(f"\033[92mDEBUG\033[0m: MFCC shape: {mfcc.shape}")
            
            # Debug image input
            print(f"\033[92mDEBUG\033[0m: Image input type: {type(image)}")
            print(f"\033[92mDEBUG\033[0m: Image shape: {image.shape if hasattr(image, 'shape') else 'No shape'}")
            
            # Process image data and handle None case
            if image is None:
                return "Error: No image provided. Please upload an image."
                
            # Handle different image formats
            if isinstance(image, np.ndarray):
                # Check if image is properly formatted (H, W, C) with 3 channels
                if len(image.shape) == 3 and image.shape[2] == 3:
                    # Convert to tensor with shape (C, H, W) as expected by PyTorch
                    img = torch.tensor(image).float().permute(2, 0, 1)
                    print(f"\033[92mDEBUG\033[0m: Converted image to tensor with shape: {img.shape}")
                elif len(image.shape) == 2:
                    # Grayscale image, expand to 3 channels
                    img = torch.tensor(image).float().unsqueeze(0).repeat(3, 1, 1)
                    print(f"\033[92mDEBUG\033[0m: Converted grayscale image to RGB tensor with shape: {img.shape}")
                else:
                    return f"Error: Unexpected image shape: {image.shape}. Expected RGB or grayscale image."
            else:
                return f"Error: Unexpected image format: {type(image)}. Expected numpy array."
            
            # Scale pixel values to [0, 1] if needed
            if img.max() > 1.0:
                img = img / 255.0
                print(f"\033[92mDEBUG\033[0m: Scaled image pixel values to range [0, 1]")
            
            # Get image dimensions and check if they're reasonable
            print(f"\033[92mDEBUG\033[0m: Final image tensor shape before processing: {img.shape}")
            
            # Process image
            try:
                img_processed = process_image_data(img)
                if img_processed is None:
                    return "Error: Failed to process image data. Make sure your image clearly shows a watermelon."
                    
                img_processed = img_processed.to(device)
                print(f"\033[92mDEBUG\033[0m: Processed image shape: {img_processed.shape}")
            except Exception as e:
                print(f"\033[91mERROR\033[0m: Image processing error: {str(e)}")
                return f"Error in image processing: {str(e)}"
            
            # Run inference
            try:
                # Based on the error, it seems infer() expects file paths, not tensors
                # Let's create temporary files for the processed data
                temp_dir = os.path.join(os.getcwd(), "temp")
                os.makedirs(temp_dir, exist_ok=True)
                
                # Save the audio to a temporary file if infer expects a file path
                temp_audio_path = os.path.join(temp_dir, "temp_audio.wav")
                if not isinstance(audio, str) and isinstance(audio, tuple) and len(audio) >= 2:
                    # If we have the original audio data and sample rate
                    audio_array = audio[-1]
                    sr = audio[0]
                    
                    # Check if the audio array is valid
                    if audio_array.size == 0:
                        return "Error: Audio data is empty. Please record a longer audio clip."
                    
                    # Get the duration of the audio
                    duration = audio_array.shape[-1] / sr
                    print(f"\033[92mDEBUG\033[0m: Audio duration: {duration:.2f} seconds")
                    
                    # Check if we have at least 1 second of audio - but don't reject, just pad if needed
                    min_duration = 1.0  # minimum 1 second of audio
                    if duration < min_duration:
                        print(f"\033[93mWARNING\033[0m: Audio is shorter than {min_duration} seconds. Padding will be applied.")
                        # Calculate samples needed to reach minimum duration
                        samples_needed = int(min_duration * sr) - audio_array.shape[-1]
                        # Pad with zeros
                        padding = np.zeros((audio_array.shape[0], samples_needed), dtype=audio_array.dtype)
                        audio_array = np.concatenate([audio_array, padding], axis=1)
                        print(f"\033[92mDEBUG\033[0m: Padded audio to shape: {audio_array.shape}")
                    
                    # Make sure audio has 2 dimensions
                    if len(audio_array.shape) == 1:
                        audio_array = np.expand_dims(audio_array, axis=0)
                    
                    print(f"\033[92mDEBUG\033[0m: Audio array shape before saving: {audio_array.shape}, sr: {sr}")
                    
                    # Make sure it's in the right format for torchaudio.save
                    audio_tensor = torch.tensor(audio_array).float()
                    if audio_tensor.dim() == 1:
                        audio_tensor = audio_tensor.unsqueeze(0)
                    
                    torchaudio.save(temp_audio_path, audio_tensor, sr)
                    print(f"\033[92mDEBUG\033[0m: Saved temporary audio file to {temp_audio_path}")
                    
                    # Let's also process the audio here to verify it works
                    test_mfcc = process_audio_data(audio_tensor, sr)
                    if test_mfcc is None:
                        return "Error: Unable to process the audio. Please try recording a different audio sample."
                    else:
                        print(f"\033[92mDEBUG\033[0m: Audio pre-check passed. MFCC shape: {test_mfcc.shape}")
                    
                    audio_path = temp_audio_path
                else:
                    # If we don't have a valid path, return an error
                    return "Error: Cannot process audio for inference. Invalid audio format."
                
                # Save the image to a temporary file if infer expects a file path
                temp_image_path = os.path.join(temp_dir, "temp_image.jpg")
                if isinstance(image, np.ndarray):
                    import cv2
                    cv2.imwrite(temp_image_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
                    print(f"\033[92mDEBUG\033[0m: Saved temporary image file to {temp_image_path}")
                    image_path = temp_image_path
                else:
                    # If we don't have a valid image, return an error
                    return "Error: Cannot process image for inference. Invalid image format."
                
                # Create a modified version of infer that handles None returns
                def safe_infer(audio_path, image_path, model, device):
                    try:
                        return infer(audio_path, image_path, model, device)
                    except Exception as e:
                        print(f"\033[91mERROR\033[0m: Error in infer function: {str(e)}")
                        # Try a more direct approach
                        try:
                            # Load audio and process
                            audio, sr = torchaudio.load(audio_path)
                            mfcc = process_audio_data(audio, sr)
                            if mfcc is None:
                                raise ValueError("Audio processing failed - MFCC is None")
                            mfcc = mfcc.to(device)
                            
                            # Load image and process
                            image = cv2.imread(image_path)
                            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                            image_tensor = torch.tensor(image).float().permute(2, 0, 1) / 255.0
                            img_processed = process_image_data(image_tensor)
                            if img_processed is None:
                                raise ValueError("Image processing failed - processed image is None")
                            img_processed = img_processed.to(device)
                            
                            # Run model inference
                            with torch.no_grad():
                                prediction = model(mfcc, img_processed)
                            return prediction
                        except Exception as e2:
                            print(f"\033[91mERROR\033[0m: Fallback inference also failed: {str(e2)}")
                            raise
                
                # Call our safer version
                print(f"\033[92mDEBUG\033[0m: Calling safe_infer with audio_path={audio_path}, image_path={image_path}")
                sweetness = safe_infer(audio_path, image_path, model, device)
                if sweetness is None:
                    return "Error: The model was unable to make a prediction. Please try with different inputs."
                    
                print(f"\033[92mDEBUG\033[0m: Inference result: {sweetness.item()}")
                return f"Predicted Sweetness: {sweetness.item():.2f}/10"
            except Exception as e:
                import traceback
                print(f"\033[91mERROR\033[0m: Inference failed: {str(e)}")
                print(f"\033[91mTraceback\033[0m: {traceback.format_exc()}")
                return f"Error during inference: {str(e)}"
            
        except Exception as e:
            import traceback
            print(f"\033[91mERROR\033[0m: Prediction failed: {str(e)}")
            print(f"\033[91mTraceback\033[0m: {traceback.format_exc()}")
            return f"Error processing input: {str(e)}"

    audio_input = gr.Audio(label="Upload or Record Audio")
    image_input = gr.Image(label="Upload or Capture Image")
    output = gr.Textbox(label="Predicted Sweetness")

    interface = gr.Interface(
        fn=predict,
        inputs=[audio_input, image_input],
        outputs=output,
        title="Watermelon Sweetness Predictor",
        description="Upload an audio file and an image to predict the sweetness of a watermelon."
    )

    try:
        interface.launch(share=True)  # Enable sharing to avoid localhost access issues
    except Exception as e:
        print(f"\033[91mERROR\033[0m: Failed to launch interface: {e}")
        print("\033[93mTIP\033[0m: If you're running in a remote environment or container, try setting additional parameters:")
        print("    interface.launch(server_name='0.0.0.0', share=True)")