Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import librosa | |
| import numpy as np | |
| import tensorflow as tf | |
| from tensorflow.keras.models import load_model | |
| import matplotlib.pyplot as plt | |
| from PIL import Image | |
| from googleapiclient.discovery import build | |
| from google.oauth2 import service_account | |
| from googleapiclient.http import MediaIoBaseUpload | |
| from flask import jsonify | |
| import logging | |
| from datetime import datetime | |
| import subprocess | |
| import tempfile | |
| import shutil | |
| from database import update_hive_health_in_db | |
| # Ensure matplotlib runs in non-GUI mode | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| # Configure logging | |
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') | |
| logger = logging.getLogger(__name__) | |
| # Google Drive API Setup | |
| SCOPES = ['https://www.googleapis.com/auth/drive.file'] | |
| PARENT_FOLDER_ID = os.getenv("GOOGLE_DRIVE_FOLDER_ID") | |
| service_account_info = json.loads(os.getenv("G_Drive_Credentials")) | |
| credentials = service_account.Credentials.from_service_account_info(service_account_info) | |
| # Initialize Google Drive API | |
| drive_service = build("drive", "v3", credentials=credentials) | |
| # Load the bee/no bee model | |
| MODEL_PATH = "./mobilenet_best_model_merged_bee_nobee.keras" | |
| bee_model = load_model(MODEL_PATH) | |
| # Load the queen model (MobileNet) | |
| QUEEN_MODEL_PATH = "./mobilenet_best_model_merged_queen_noqueen.keras" | |
| queen_model = load_model(QUEEN_MODEL_PATH) | |
| # Load the mite attack model (MobileNet) | |
| MITE_MODEL_PATH = "./mobilenet_best_model_merged_mite_nomite.keras" | |
| mite_model = load_model(MITE_MODEL_PATH) | |
| def check_ffmpeg(): | |
| """Check if FFmpeg is available.""" | |
| return shutil.which("ffmpeg") is not None | |
| def convert_to_wav(input_path, output_path): | |
| """Convert audio file to WAV format using FFmpeg.""" | |
| try: | |
| if not os.path.exists(input_path): | |
| raise FileNotFoundError("Input audio file does not exist") | |
| if not check_ffmpeg(): | |
| raise RuntimeError("FFmpeg is not installed or not found in PATH") | |
| cmd = [ | |
| "ffmpeg", | |
| "-i", input_path, | |
| "-acodec", "pcm_s16le", | |
| "-ar", "44100", | |
| "-ac", "2", | |
| "-y", | |
| output_path | |
| ] | |
| result = subprocess.run( | |
| cmd, | |
| stdout=subprocess.PIPE, | |
| stderr=subprocess.PIPE, | |
| text=True | |
| ) | |
| if result.returncode != 0: | |
| raise RuntimeError(f"FFmpeg conversion failed: {result.stderr}") | |
| logger.info(f"Converted {input_path} to {output_path}") | |
| return output_path | |
| except Exception as e: | |
| logger.error(f"Error converting audio to WAV: {str(e)}") | |
| raise | |
| def get_or_create_folder(folder_name, parent_id): | |
| """Retrieve or create a folder in Google Drive and return its ID.""" | |
| query = f"name='{folder_name}' and mimeType='application/vnd.google-apps.folder' and '{parent_id}' in parents and trashed=false" | |
| response = drive_service.files().list(q=query, spaces='drive', fields='files(id, name)').execute() | |
| files = response.get('files', []) | |
| if files: | |
| return files[0].get('id') | |
| folder_metadata = { | |
| 'name': folder_name, | |
| 'mimeType': 'application/vnd.google-apps.folder', | |
| 'parents': [parent_id] | |
| } | |
| folder = drive_service.files().create(body=folder_metadata, fields='id').execute() | |
| return folder.get('id') | |
| def upload_to_drive(audio_file, result, user_predict=None): | |
| """Uploads an audio file to the appropriate Google Drive folder(s) based on prediction result and user_predict.""" | |
| try: | |
| file_ids = [] | |
| if result == "not bee": | |
| folder_name = "not bee" | |
| folder_id = get_or_create_folder(folder_name, PARENT_FOLDER_ID) | |
| file_metadata = { | |
| 'name': audio_file.filename, | |
| 'parents': [folder_id] | |
| } | |
| media = MediaIoBaseUpload(audio_file, mimetype=audio_file.content_type, resumable=True) | |
| file = drive_service.files().create( | |
| body=file_metadata, | |
| media_body=media, | |
| fields='id' | |
| ).execute() | |
| file_ids.append(file.get('id')) | |
| logger.info(f"Uploaded {audio_file.filename} to Google Drive folder '{folder_name}' with file ID: {file.get('id')}") | |
| else: | |
| bee_folder_id = get_or_create_folder("bee", PARENT_FOLDER_ID) | |
| audio_file.seek(0) | |
| file_metadata = { | |
| 'name': audio_file.filename, | |
| 'parents': [bee_folder_id] | |
| } | |
| media = MediaIoBaseUpload(audio_file, mimetype=audio_file.content_type, resumable=True) | |
| file = drive_service.files().create( | |
| body=file_metadata, | |
| media_body=media, | |
| fields='id' | |
| ).execute() | |
| file_ids.append(file.get('id')) | |
| logger.info(f"Uploaded {audio_file.filename} to Google Drive folder 'bee' with file ID: {file.get('id')}") | |
| if user_predict: | |
| user_predict = user_predict.strip().lower() | |
| valid_folders = {'healthy', 'no queen', 'mite attack', 'chalkbrood'} | |
| if user_predict in valid_folders: | |
| folder_id = get_or_create_folder(user_predict, PARENT_FOLDER_ID) | |
| audio_file.seek(0) | |
| file_metadata = { | |
| 'name': audio_file.filename, | |
| 'parents': [folder_id] | |
| } | |
| media = MediaIoBaseUpload(audio_file, mimetype=audio_file.content_type, resumable=True) | |
| file = drive_service.files().create( | |
| body=file_metadata, | |
| media_body=media, | |
| fields='id' | |
| ).execute() | |
| file_ids.append(file.get('id')) | |
| logger.info(f"Uploaded {audio_file.filename} to Google Drive folder '{user_predict}' with file ID: {file.get('id')}") | |
| else: | |
| logger.warning(f"Ignoring invalid user_predict value: {user_predict}") | |
| return file_ids[0] | |
| except Exception as e: | |
| logger.error(f"Error uploading to Google Drive: {str(e)}") | |
| raise | |
| def create_mel_spectrogram(audio_segment, sr): | |
| """Creates a mel spectrogram from an audio segment.""" | |
| try: | |
| spectrogram = librosa.feature.melspectrogram(y=audio_segment, sr=sr, n_mels=128) | |
| spectrogram_db = librosa.power_to_db(spectrogram, ref=np.max) | |
| plt.figure(figsize=(2, 2), dpi=100) | |
| plt.axis('off') | |
| plt.imshow(spectrogram_db, aspect='auto', cmap='magma', origin='lower') | |
| plt.tight_layout(pad=0) | |
| temp_image_path = "/tmp/temp_spectrogram.png" | |
| plt.savefig(temp_image_path, bbox_inches='tight', pad_inches=0) | |
| plt.close() | |
| img = Image.open(temp_image_path).convert('RGB') | |
| img = img.resize((224, 224)) | |
| img_array = np.array(img) / 255.0 | |
| os.remove(temp_image_path) | |
| return img_array | |
| except Exception as e: | |
| logger.error(f"Error creating spectrogram: {e}") | |
| return None | |
| def predict_queen_audio(file_path, model): | |
| """Processes audio for queen detection using 10-second segments.""" | |
| try: | |
| y, sr = librosa.load(file_path, sr=None) | |
| duration = librosa.get_duration(y=y, sr=sr) | |
| if duration <= 10: | |
| return {"error": "Audio file must be longer than 10 seconds"} | |
| healthy_count = 0 | |
| total_segments = 0 | |
| segment_start = 0 | |
| while segment_start < duration: | |
| segment_end = min(segment_start + 10, duration) | |
| if segment_end - segment_start < 10 and segment_start > 0: | |
| segment_start = max(0, duration - 10) | |
| segment_end = duration | |
| audio_segment = y[int(segment_start * sr):int(segment_end * sr)] | |
| spectrogram = create_mel_spectrogram(audio_segment, sr) | |
| if spectrogram is not None: | |
| spectrogram = np.expand_dims(spectrogram, axis=0) | |
| prediction = model.predict(spectrogram) | |
| probability = prediction[0][0] # Assuming binary classification (0: no queen, 1: healthy) | |
| if probability >= 0.8: # Aligned with predict_audio for consistency | |
| healthy_count += 1 | |
| total_segments += 1 | |
| segment_start += 10 | |
| if total_segments > 0: | |
| healthy_percentage = (healthy_count / total_segments) * 100 | |
| result = "healthy" if healthy_percentage >= 70 else "no queen" | |
| return result | |
| else: | |
| return {"error": "No valid segments processed"} | |
| except Exception as e: | |
| logger.error(f"Error in queen prediction: {e}") | |
| return {"error": str(e)} | |
| def predict_mite_audio(file_path, model): | |
| """Processes audio for mite attack detection using 10-second segments.""" | |
| try: | |
| y, sr = librosa.load(file_path, sr=None) | |
| duration = librosa.get_duration(y=y, sr=sr) | |
| if duration <= 10: | |
| return {"error": "Audio file must be longer than 10 seconds"} | |
| healthy_count = 0 | |
| total_segments = 0 | |
| segment_start = 0 | |
| while segment_start < duration: | |
| segment_end = min(segment_start + 10, duration) | |
| if segment_end - segment_start < 10 and segment_start > 0: | |
| segment_start = max(0, duration - 10) | |
| segment_end = duration | |
| audio_segment = y[int(segment_start * sr):int(segment_end * sr)] | |
| spectrogram = create_mel_spectrogram(audio_segment, sr) | |
| if spectrogram is not None: | |
| spectrogram = np.expand_dims(spectrogram, axis=0) | |
| prediction = model.predict(spectrogram) | |
| probability = prediction[0][0] # Assuming binary classification (0: mite attack, 1: healthy) | |
| if probability >= 0.8: # Consistent threshold | |
| healthy_count += 1 | |
| total_segments += 1 | |
| segment_start += 10 | |
| if total_segments > 0: | |
| healthy_percentage = (healthy_count / total_segments) * 100 | |
| result = "healthy" if healthy_percentage >= 70 else "mite attack" | |
| return result | |
| else: | |
| return {"error": "No valid segments processed"} | |
| except Exception as e: | |
| logger.error(f"Error in mite attack prediction: {e}") | |
| return {"error": str(e)} | |
| def predict_audio(audio_path, request_id): | |
| """Predicts whether an audio file contains bee sounds.""" | |
| try: | |
| y, sr = librosa.load(audio_path, sr=None) | |
| if y is None or sr is None: | |
| return {"error": "Failed to load audio"} | |
| duration = librosa.get_duration(y=y, sr=sr) | |
| if duration <= 10: | |
| return {"error": "Audio file must be longer than 10 seconds"} | |
| bee_count = 0 | |
| total_segments = 0 | |
| segment_start = 0 | |
| while segment_start < duration: | |
| segment_end = min(segment_start + 10, duration) | |
| if segment_end - segment_start < 10 and segment_start > 0: | |
| segment_start = max(0, duration - 10) | |
| segment_end = duration | |
| audio_segment = y[int(segment_start * sr):int(segment_end * sr)] | |
| spectrogram = create_mel_spectrogram(audio_segment, sr) | |
| if spectrogram is not None: | |
| spectrogram = np.expand_dims(spectrogram, axis=0) | |
| prediction = bee_model.predict(spectrogram) | |
| probability = prediction[0][0] | |
| if probability <= 0.2: | |
| bee_count += 1 | |
| total_segments += 1 | |
| segment_start += 10 | |
| if total_segments > 0: | |
| bee_percentage = (bee_count / total_segments) * 100 | |
| result = "bee" if bee_percentage >= 70 else "not bee" | |
| logger.info(f"Request {request_id} - Prediction result: {result}") | |
| return {"result": result} | |
| else: | |
| return {"result": "try again"} | |
| except Exception as e: | |
| logger.error(f"Request {request_id} - Error during bee prediction: {e}") | |
| return {"error": str(e)} | |
| def handle_predict(request, save_prediction): | |
| """Handles the prediction logic for the /predict route.""" | |
| request_id = datetime.now().strftime("%Y%m%d_%H%M%S_%f") | |
| logger.info(f"Request {request_id} - Starting prediction") | |
| if 'audio' not in request.files or 'user_id' not in request.form: | |
| logger.error(f"Request {request_id} - Missing data") | |
| return jsonify({"error": "Missing data"}), 400 | |
| user_id = request.form['user_id'] | |
| audio_file = request.files['audio'] | |
| hive_id = request.form.get('hive_id') | |
| user_predict = request.form.get('user_predict') | |
| if audio_file.filename == '': | |
| logger.error(f"Request {request_id} - No file selected") | |
| return jsonify({"error": "No file selected"}), 400 | |
| if not check_ffmpeg(): | |
| logger.error(f"Request {request_id} - FFmpeg not found") | |
| return jsonify({"error": "Server error: FFmpeg is not available"}), 500 | |
| temp_dir = tempfile.gettempdir() | |
| original_filename = audio_file.filename | |
| original_ext = os.path.splitext(original_filename)[1].lower() | |
| temp_original_path = os.path.join(temp_dir, f"temp_audio_{request_id}{original_ext}") | |
| audio_file.save(temp_original_path) | |
| temp_wav_path = os.path.join(temp_dir, f"temp_audio_{request_id}.wav") | |
| if original_ext not in ['.wav', '.mp3']: | |
| try: | |
| convert_to_wav(temp_original_path, temp_wav_path) | |
| except Exception as e: | |
| os.remove(temp_original_path) | |
| logger.error(f"Request {request_id} - Audio conversion failed: {str(e)}") | |
| return jsonify({"error": f"Failed to process audio file: {str(e)}"}), 400 | |
| else: | |
| temp_wav_path = temp_original_path | |
| logger.info(f"Request {request_id} - Using original file (no conversion needed): {temp_wav_path}") | |
| prediction_result = predict_audio(temp_wav_path, request_id) | |
| if "error" in prediction_result: | |
| os.remove(temp_original_path) | |
| if temp_wav_path != temp_original_path and os.path.exists(temp_wav_path): | |
| os.remove(temp_wav_path) | |
| logger.error(f"Request {request_id} - Prediction failed: {prediction_result['error']}") | |
| return jsonify({"result": "try again"}), 400 | |
| result = prediction_result["result"].lower() | |
| if result == "try again": | |
| os.remove(temp_original_path) | |
| if temp_wav_path != temp_original_path and os.path.exists(temp_wav_path): | |
| os.remove(temp_wav_path) | |
| logger.info(f"Request {request_id} - Result: try again") | |
| return jsonify({"result": "try again"}) | |
| if result == "bee": | |
| queen_result = predict_queen_audio(temp_wav_path, queen_model) | |
| mite_result = predict_mite_audio(temp_wav_path, mite_model) | |
| if "error" in queen_result or "error" in mite_result: | |
| os.remove(temp_original_path) | |
| if temp_wav_path != temp_original_path and os.path.exists(temp_wav_path): | |
| os.remove(temp_wav_path) | |
| logger.error(f"Request {request_id} - Queen prediction failed: {queen_result.get('error', 'Unknown')}, Mite prediction failed: {mite_result.get('error', 'Unknown')}") | |
| return jsonify({"result": "try again"}), 400 | |
| logger.info(f"Request {request_id} - Queen prediction result: {queen_result}, Mite prediction result: {mite_result}") | |
| # Combine results according to specified logic | |
| if queen_result == "healthy" and mite_result == "healthy": | |
| result = "healthy" | |
| elif queen_result == "no queen" and mite_result == "healthy": | |
| result = "no queen" | |
| elif queen_result == "healthy" and mite_result == "mite attack": | |
| result = "mite attack" | |
| elif queen_result == "no queen" and mite_result == "mite attack": | |
| result = "no queen,mite attack" | |
| else: | |
| result = "try again" # Fallback for unexpected cases | |
| logger.warning(f"Request {request_id} - Unexpected combination: queen={queen_result}, mite={mite_result}") | |
| if user_predict and hive_id: | |
| try: | |
| user_predict = user_predict.strip().lower() | |
| update_hive_health_in_db(hive_id, user_predict) | |
| logger.info(f"Request {request_id} - Updated hive {hive_id} health_status to {user_predict}") | |
| except Exception as e: | |
| logger.error(f"Request {request_id} - Failed to update hive health status: {str(e)}") | |
| return jsonify({"error": f"Failed to update hive health status: {str(e)}"}), 400 | |
| file_id = None | |
| try: | |
| audio_file.seek(0) | |
| file_id = upload_to_drive(audio_file, result, user_predict) | |
| except Exception as e: | |
| logger.error(f"Request {request_id} - Failed to upload to Google Drive: {str(e)}") | |
| # Continue with saving prediction and returning result with file_id=None | |
| save_prediction(user_id, audio_file.filename, result, file_id, hive_id, user_predict) | |
| os.remove(temp_original_path) | |
| if temp_wav_path != temp_original_path and os.path.exists(temp_wav_path): | |
| os.remove(temp_wav_path) | |
| logger.info(f"Request {request_id} - Final result: {result}, file_id: {file_id}") | |
| return jsonify({"result": result, "file_id": file_id}) |