Update app.py

app.py

@@ -5,18 +5,16 @@ import logging
 from flask import Flask, request, jsonify, render_template, send_from_directory, url_for
 from werkzeug.utils import secure_filename
 from ultralytics import YOLO
-import google.generativeai as genai
+import google.generativeai as genai  # Import Gemini API
 import cv2
 from huggingface_hub import hf_hub_download
-import io
-import tempfile
 
 # --- Basic Setup & Configuration ---
 logging.basicConfig(level=logging.INFO)
 app = Flask(__name__)
-app.config['UPLOAD_FOLDER'] = 'uploads'  # Save uploads in the persistent uploads folder
-app.config['RESULT_FOLDER'] = 'static/results'  # Save result images in static/results
-app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16 MB file upload limit
+app.config['UPLOAD_FOLDER'] = 'uploads'
+app.config['RESULT_FOLDER'] = os.path.join('static', 'results')
+app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
 ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'webp'}
 
 # --- Ensure Folders Exist ---
@@ -42,35 +40,63 @@ if not bird_data:
     logging.warning("Bird data is empty. Features relying on it might not work.")
 
 # --- Load YOLOv8 Model ---
-model = None
+# --- Download and Load YOLOv8 Model ---
+model = None # Initialize as None
 try:
-    HF_REPO_ID = "aevnum/avian-intelligence-yolo" 
+    # Define Hugging Face repo details - CHANGE THESE
+    HF_REPO_ID = "aevnum/avian-intelligence-yolo" # <<<--- YOUR HF REPO ID
     HF_FILENAME = "best.pt"
-    MODEL_CACHE_DIR = "/tmp/model_cache"
+    MODEL_CACHE_DIR = "model_cache" # Can be any directory name
 
+    # Ensure the local model cache directory exists
     os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
 
     logging.info(f"Downloading model {HF_FILENAME} from {HF_REPO_ID}...")
+    # Use HF Token from environment variable for download
     hf_token = os.environ.get('HUGGING_FACE_HUB_TOKEN')
     if not hf_token:
-        logging.warning("HUGGING_FACE_HUB_TOKEN not set. Download might fail for private repos.")
+         logging.warning("HUGGING_FACE_HUB_TOKEN not set. Download might fail for private repos or hit rate limits.")
 
     downloaded_model_path = hf_hub_download(
         repo_id=HF_REPO_ID,
         filename=HF_FILENAME,
         cache_dir=MODEL_CACHE_DIR,
-        force_filename=HF_FILENAME,
+        force_filename=HF_FILENAME, # Helps ensure consistent naming if cache is used
         token=hf_token
     )
     logging.info(f"Model downloaded to: {downloaded_model_path}")
 
+    # Load the downloaded model
     model = YOLO(downloaded_model_path)
     logging.info("YOLOv8 model loaded successfully from downloaded file.")
 
 except Exception as e:
-    logging.exception("Error downloading or loading YOLOv8 model from Hugging Face Hub") 
+    logging.exception("Error downloading or loading YOLOv8 model from Hugging Face Hub") # Log full traceback
     model = None
 
+# --- Configure Gemini API ---
+model_gemini = None # Initialize as None
+try:
+    gemini_api_key = os.environ.get('GEMINI_API_KEY') # Get key from environment
+    if not gemini_api_key:
+        logging.warning("GEMINI_API_KEY environment variable not set. Chat feature will be disabled.")
+    else:
+        logging.info("Configuring Gemini API...")
+        genai.configure(api_key=gemini_api_key)
+        # Consider making model name configurable too via env var if needed
+        # GEMINI_MODEL = os.environ.get('GEMINI_MODEL_NAME', 'gemini-1.5-flash-latest')
+        model_gemini = genai.GenerativeModel('gemini-2.0-flash') # Or use variable
+        logging.info(f"Gemini client configured with model.")
+
+except Exception as e:
+    logging.exception("Failed to initialize Gemini client") # Log full traceback
+    model_gemini = None
+
+# --- Helper Functions ---
+def allowed_file(filename):
+    return '.' in filename and \
+           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
+
 # --- Prediction Function ---
 def predict_birds(image_path):
     if not model:
@@ -85,11 +111,10 @@ def predict_birds(image_path):
 
         processed_results = results[0]
 
-        # Save the processed image to the persistent folder (static/results)
-        result_filename = f"prediction_{uuid.uuid4().hex}.jpg"
-        result_image_path = os.path.join(app.config['RESULT_FOLDER'], result_filename)
-        processed_results.save(result_image_path)
-        logging.info(f"Saved prediction result image to {result_image_path}")
+        output_filename = f"{uuid.uuid4()}.jpg"
+        output_path = os.path.join(app.config['RESULT_FOLDER'], output_filename)
+        processed_results.save(filename=output_path)
+        logging.info(f"Saved prediction result image to {output_path}")
 
         detected_classes = []
         names = processed_results.names
@@ -110,7 +135,8 @@ def predict_birds(image_path):
 
         detected_classes.sort(key=lambda x: x['confidence'], reverse=True)
 
-        return result_image_path, None, detected_classes
+        relative_output_path = f"results/{output_filename}"
+        return relative_output_path, None, detected_classes
 
     except Exception as e:
         logging.exception(f"Error during prediction for {image_path}")
@@ -136,34 +162,45 @@ def handle_prediction():
 
     if file and allowed_file(file.filename):
         filename = secure_filename(file.filename)
-        # Save the uploaded file to persistent uploads folder
-        upload_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
-        file.save(upload_path)
-        logging.info(f"Uploaded file saved to {upload_path}")
-
-        result_image_path, error_msg, detected_classes = predict_birds(upload_path)
-
-        logging.info(f"predict_birds returned image path: {result_image_path}")
-
-        if error_msg:
-            logging.error(f"Prediction error message: {error_msg}")
-            return jsonify({'error': error_msg}), 500
-
-        if result_image_path:
-            # Construct result image URL for the static folder
-            result_image_url = url_for('static', filename=f"results/{os.path.basename(result_image_path)}", _external=False)
-            logging.info(f"Generated result_image_url: {result_image_url}")
-        else:
-            result_image_url = None
-            logging.warning("No result_image_path returned, URL will be null.")
-
-        response_data = {
-            'result_image_url': result_image_url,
-            'detections': detected_classes
-        }
-        logging.info(f"Returning JSON: {response_data}")
-        return jsonify(response_data)
-
+        temp_filename = f"{uuid.uuid4()}_{filename}"
+        temp_filepath = os.path.join(app.config['UPLOAD_FOLDER'], temp_filename)
+
+        try:
+            file.save(temp_filepath)
+            logging.info(f"Uploaded file saved temporarily to {temp_filepath}")
+
+            result_image_rel_path, error_msg, detected_classes = predict_birds(temp_filepath)
+
+            logging.info(f"predict_birds returned relative path: {result_image_rel_path}")
+
+            if error_msg:
+                logging.error(f"Prediction error message: {error_msg}")
+                return jsonify({'error': error_msg}), 500
+
+            if result_image_rel_path:
+                result_image_url = url_for('static', filename=result_image_rel_path, _external=False)
+                logging.info(f"Generated result_image_url: {result_image_url}")
+            else:
+                result_image_url = None
+                logging.warning("No result_image_rel_path returned, URL will be null.")
+
+            response_data = {
+                'result_image_url': result_image_url,
+                'detections': detected_classes
+            }
+            logging.info(f"Returning JSON: {response_data}")
+            return jsonify(response_data)
+
+        except Exception as e:
+            logging.exception("Error handling prediction request")
+            return jsonify({'error': 'Failed to process image'}), 500
+        finally:
+            if os.path.exists(temp_filepath):
+                try:
+                    os.remove(temp_filepath)
+                    logging.info(f"Removed temporary file: {temp_filepath}")
+                except OSError as e:
+                    logging.error(f"Error removing temporary file {temp_filepath}: {e}")
     else:
         return jsonify({'error': 'Invalid file type'}), 400
 
@@ -183,32 +220,47 @@ def get_bird_info(bird_name):
         logging.warning(f"Bird info requested for '{safe_bird_name}', but not found in data.")
         return jsonify({'error': 'Bird species not found in database'}), 404
 
-# --- Cleanup after serving ---
-# @app.after_request
-# def cleanup(response):
-#     # Clean up the uploaded file after serving
-#     uploaded_files = os.listdir(app.config['UPLOAD_FOLDER'])
-#     for filename in uploaded_files:
-#         file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
-#         try:
-#             os.remove(file_path)
-#             logging.info(f"Removed uploaded file: {file_path}")
-#         except OSError as e:
-#             logging.error(f"Error removing uploaded file {file_path}: {e}")
-
-#     # Clean up result images after serving
-#     result_files = os.listdir(app.config['RESULT_FOLDER'])
-#     for filename in result_files:
-#         file_path = os.path.join(app.config['RESULT_FOLDER'], filename)
-#         try:
-#             os.remove(file_path)
-#             logging.info(f"Removed result file: {file_path}")
-#         except OSError as e:
-#             logging.error(f"Error removing result file {file_path}: {e}")
-
-#     return response
+@app.route('/chat', methods=['POST'])
+def handle_chat():
+    if not model_gemini:
+        return jsonify({'reply': "Sorry, the chat feature is not configured or the API key is missing."}), 503
+
+    data = request.get_json()
+    if not data or 'bird_name' not in data or 'message' not in data:
+        return jsonify({'error': 'Missing bird_name or message in request'}), 400
+
+    bird_name = data['bird_name']
+    user_message = data['message']
+    chat_history = data.get('history', [])
+
+    bird_details = bird_data.get(bird_name, {})
+    context_summary = f"Genus: {bird_details.get('genus', 'N/A')}, Locations: {bird_details.get('locations', 'N/A')}, Info: {bird_details.get('short_info', 'N/A')}."
+
+    messages = [
+        {"role": "system", "content": f"You are a helpful ornithology assistant specializing in bird information. The user is asking about the '{bird_name}'. Basic info: {context_summary}. Keep answers concise and relevant to birds."},
+    ]
+    for entry in chat_history[-4:]:
+        messages.append({"role": entry["role"], "content": entry["content"]})
+
+    messages.append({"role": "user", "content": user_message})
+
+    try:
+        logging.info(f"Sending request to Gemini for bird: {bird_name}")
+        # Gemini API interaction
+        prompt = "\n".join([msg["content"] for msg in messages]) #convert messages to one string.
+        response = model_gemini.generate_content(prompt)
+        ai_reply = response.text.strip()
+        logging.info(f"Received reply from Gemini for bird: {bird_name}")
+
+        return jsonify({'reply': ai_reply})
+
+    except Exception as e:
+        logging.exception("Unexpected error in chat handler")
+        return jsonify({'reply': "Sorry, an unexpected error occurred while contacting the AI assistant."}), 500
 
 # --- Main Execution ---
 if __name__ == '__main__':
-    port = int(os.environ.get('PORT', 5000))  # Default to 5000 if PORT not set
+    # Use host='0.0.0.0' to be accessible within the container
+    # Port is usually set by the deployment platform via PORT env var
+    port = int(os.environ.get('PORT', 5000)) # Default to 5000 if PORT not set
     app.run(host='0.0.0.0', port=port)