app.py

from flask import Flask, render_template, request, jsonify
from werkzeug.utils import secure_filename
import os
from PIL import Image
import pytesseract
from transformers import BlipProcessor, BlipForConditionalGeneration
import torch
import torchvision
from torchvision import transforms
from torchvision.models.detection import FasterRCNN_ResNet50_FPN_Weights
import torch.nn.functional as F
import subprocess
import time 
import warnings
import logging
import json
from datetime import datetime, timedelta
import threading
import gc
import shutil
from pathlib import Path
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
from functools import wraps
from concurrent.futures import ThreadPoolExecutor
import asyncio
from functools import lru_cache
import base64
from io import BytesIO
import requests
import ollama

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('app.log'),
        logging.StreamHandler()
    ]
)

# Set Tesseract path
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size
app.config['MAX_CHAT_HISTORY'] = 100  # Maximum number of chats to store
app.config['CHAT_EXPIRY_DAYS'] = 7  # Days after which chats are considered expired

# GRAQ API Key
GRAQ_API_KEY = "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"

# Ensure upload directory exists
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# Initialize models
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")

detection_model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
    weights=FasterRCNN_ResNet50_FPN_Weights.DEFAULT
).to(device)
detection_model.eval()

# Load BLIP models with force_download=False
processor = BlipProcessor.from_pretrained(
    "Salesforce/blip-image-captioning-base",
    force_download=False
)
caption_model = BlipForConditionalGeneration.from_pretrained(
    "Salesforce/blip-image-captioning-base",
    force_download=False
).to(device)

# COCO class names
COCO_CLASSES = [
    '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
    'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A',
    'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
    'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
    'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
    'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table',
    'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
    'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book',
    'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]

# Global variables with cleanup
processed_image_data = {}
chat_timestamps = {}

# Initialize rate limiter
limiter = Limiter(
    app=app,
    key_func=get_remote_address,
    default_limits=["200 per day", "50 per hour"]
)

# Initialize thread pool for background tasks
executor = ThreadPoolExecutor(max_workers=4)

# Cache for processed images
@lru_cache(maxsize=100)
def get_cached_image_data(chat_id):
    return processed_image_data.get(chat_id)

# Security headers
@app.after_request
def add_security_headers(response):
    response.headers['X-Content-Type-Options'] = 'nosniff'
    response.headers['X-Frame-Options'] = 'SAMEORIGIN'
    response.headers['X-XSS-Protection'] = '1; mode=block'
    response.headers['Strict-Transport-Security'] = 'max-age=31536000; includeSubDomains'
    return response

# Rate limiting decorator
def rate_limit():
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            try:
                return f(*args, **kwargs)
            except Exception as e:
                logging.error(f"Rate limit exceeded: {str(e)}")
                return jsonify({'error': 'Rate limit exceeded. Please try again later.'}), 429
        return decorated_function
    return decorator

# File validation
def validate_file(file):
    ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
    MAX_FILE_SIZE = 16 * 1024 * 1024  # 16MB
    
    if not file:
        raise ValueError('No file provided')
        
    if not file.filename:
        raise ValueError('No filename provided')
        
    if not '.' in file.filename:
        raise ValueError('No file extension')
        
    ext = file.filename.rsplit('.', 1)[1].lower()
    if ext not in ALLOWED_EXTENSIONS:
        raise ValueError(f'File type not allowed. Allowed types: {", ".join(ALLOWED_EXTENSIONS)}')
        
    if len(file.read()) > MAX_FILE_SIZE:
        raise ValueError('File size exceeds 16MB limit')
        
    file.seek(0)  # Reset file pointer
    return True

# Configure chat storage
CHAT_STORAGE_FILE = 'chat_history.json'

def load_chat_history():
    """Load chat history from file"""
    try:
        if os.path.exists(CHAT_STORAGE_FILE):
            with open(CHAT_STORAGE_FILE, 'r') as f:
                data = json.load(f)
                return data.get('processed_image_data', {}), data.get('chat_timestamps', {})
        return {}, {}
    except Exception as e:
        logging.error(f"Error loading chat history: {str(e)}")
        return {}, {}

def save_chat_history():
    """Save chat history to file"""
    try:
        # Convert datetime objects to strings for JSON serialization
        chat_times = {
            chat_id: timestamp.isoformat() 
            for chat_id, timestamp in chat_timestamps.items()
        }
        
        data = {
            'processed_image_data': processed_image_data,
            'chat_timestamps': chat_times
        }
        
        with open(CHAT_STORAGE_FILE, 'w') as f:
            json.dump(data, f)
    except Exception as e:
        logging.error(f"Error saving chat history: {str(e)}")

# Load existing chat history on startup
processed_image_data, loaded_timestamps = load_chat_history()
chat_timestamps = {
    chat_id: datetime.fromisoformat(timestamp_str)
    for chat_id, timestamp_str in loaded_timestamps.items()
}

def cleanup_old_chats():
    """Remove old chats and free up memory"""
    current_time = datetime.now()
    expired_chats = []
    
    for chat_id, timestamp in chat_timestamps.items():
        if current_time - timestamp > timedelta(days=app.config['CHAT_EXPIRY_DAYS']):
            expired_chats.append(chat_id)
    
    for chat_id in expired_chats:
        if chat_id in processed_image_data:
            del processed_image_data[chat_id]
        if chat_id in chat_timestamps:
            del chat_timestamps[chat_id]
    
    if expired_chats:
        logging.info(f"Cleaned up {len(expired_chats)} expired chats")
        # Save changes after cleanup
        save_chat_history()
        gc.collect()  # Force garbage collection

def manage_chat_storage():
    """Ensure we don't exceed maximum chat storage"""
    if len(processed_image_data) > app.config['MAX_CHAT_HISTORY']:
        # Remove oldest chats
        sorted_chats = sorted(chat_timestamps.items(), key=lambda x: x[1])
        chats_to_remove = len(processed_image_data) - app.config['MAX_CHAT_HISTORY']
        
        for chat_id, _ in sorted_chats[:chats_to_remove]:
            if chat_id in processed_image_data:
                del processed_image_data[chat_id]
            if chat_id in chat_timestamps:
                del chat_timestamps[chat_id]
        
        logging.info(f"Removed {chats_to_remove} oldest chats to maintain storage limit")
        # Save changes after managing storage
        save_chat_history()
        gc.collect()

def start_cleanup_thread():
    """Start background thread for periodic cleanup"""
    def cleanup_loop():
        while True:
            cleanup_old_chats()
            manage_chat_storage()
            time.sleep(3600)  # Run every hour
    
    thread = threading.Thread(target=cleanup_loop, daemon=True)
    thread.start()
    return thread

# Start cleanup thread
cleanup_thread = start_cleanup_thread()

# Background task for image processing
async def process_image_async(image_path, chat_id):
    loop = asyncio.get_event_loop()
    try:
        # Process image in background thread
        results = await loop.run_in_executor(executor, process_image, image_path)
        
        # Store results
        processed_image_data[chat_id] = results
        chat_timestamps[chat_id] = datetime.now()
        
        # Save chat history immediately after processing
        save_chat_history()
        
        # Clean up if needed
        manage_chat_storage()
        
        return results
    except Exception as e:
        logging.error(f"Error in background image processing: {str(e)}")
        raise

def detect_objects(image_path):
    """Detect objects in the image using Faster R-CNN"""
    try:
        # Open image
        image = Image.open(image_path).convert('RGB')
        
        # Transform image
        transform = transforms.Compose([
            transforms.ToTensor(),
        ])
        img_tensor = transform(image).to(device)
        
        # Run detection
        with torch.no_grad():
            predictions = detection_model([img_tensor])
        
        # Process predictions
        objects = []
        for score, label in zip(predictions[0]['scores'], predictions[0]['labels']):
            if score > 0.5:  # Confidence threshold
                class_name = COCO_CLASSES[label.item()]
                confidence = score.item()
                objects.append({
                    'class': class_name,
                    'confidence': confidence
                })
        
        return objects
    except Exception as e:
        logging.error(f"Error in object detection: {str(e)}")
        return []

def extract_text(image_path):
    """Extract text from image using Tesseract OCR"""
    try:
        image = Image.open(image_path)
        text = pytesseract.image_to_string(image)
        return text.strip()
    except Exception as e:
        logging.error(f"Error in text extraction: {str(e)}")
        return ""

def generate_caption(image_path):
    """Generate image caption using BLIP model"""
    try:
        image = Image.open(image_path).convert('RGB')
        inputs = processor(image, return_tensors="pt").to(device)
        out = caption_model.generate(**inputs)
        caption = processor.decode(out[0], skip_special_tokens=True)
        return caption
    except Exception as e:
        logging.error(f"Error in caption generation: {str(e)}")
        return ""

def process_image(image_path):
    """Process image and return analysis results"""
    start_time = time.time()
    optimized_path = None
    
    try:
        # Open and resize image if too large
        with Image.open(image_path) as img:
            # Calculate new dimensions while maintaining aspect ratio
            max_size = (1024, 1024)
            img.thumbnail(max_size, Image.Resampling.LANCZOS)
            
            # Convert to RGB if necessary
            if (img.mode != 'RGB'):
                img = img.convert('RGB')
            
            # Save optimized image in the same directory as original
            optimized_path = os.path.join(
                os.path.dirname(image_path),
                f"{os.path.splitext(os.path.basename(image_path))[0]}_optimized.jpg"
            )
            img.save(optimized_path, 'JPEG', quality=85, optimize=True)
            
            # Process the optimized image
            results = {
                'objects': detect_objects(optimized_path),
                'text': extract_text(optimized_path),
                'caption': generate_caption(optimized_path),
                'image_path': image_path  # Store original path
            }
            
            processing_time = time.time() - start_time
            logging.info(f"Image processed in {processing_time:.2f} seconds")
            
            # Save chat history after processing new image
            save_chat_history()
            
            return results
            
    except Exception as e:
        logging.error(f"Error processing image: {str(e)}")
        raise
    finally:
        # Clean up optimized image after processing
        if optimized_path and os.path.exists(optimized_path):
            try:
                os.remove(optimized_path)
            except Exception as e:
                logging.error(f"Error cleaning up optimized image: {str(e)}")

def ensure_ollama_model(model_name):
    """Check if the specified Ollama model is available and pull it if not"""
    try:
        # List available models
        available_models = ollama.list()
        model_names = [model['name'] for model in available_models.get('models', [])]
        
        if model_name not in model_names:
            logging.info(f"Pulling Ollama model {model_name}...")
            ollama.pull(model_name)
            logging.info(f"Successfully pulled {model_name}")
        else:
            logging.info(f"Ollama model {model_name} is already available")
    except Exception as e:
        logging.error(f"Error ensuring Ollama model {model_name}: {str(e)}")
        raise Exception("Failed to ensure Ollama model availability. Please ensure Ollama is installed and running.")

def generate_llm_response(question, image_data):
    """Generate response using Ollama"""
    try:
        prompt = f"""Analyze this image and answer the question professionally:

Image Analysis:
- Caption: {image_data['caption']}
- Detected Objects: {', '.join([obj['class'] for obj in image_data['objects']])}
- Extracted Text: {image_data['text'] if image_data['text'] else 'No text detected'}

Question: {question}"""

        try:
            # Use Ollama with the specified model
            response = ollama.chat(
                model='deepseek-r1:latest',
                messages=[{
                    'role': 'system',
                    'content': 'You are a professional image analysis assistant. Provide clear, accurate answers about images.'
                }, {
                    'role': 'user',
                    'content': prompt
                }],
                options={
                    'temperature': 0.7,
                    'num_ctx': 4096
                }
            )
            
            if response and 'message' in response and 'content' in response['message']:
                return response['message']['content']
            else:
                raise Exception("Invalid response format from Ollama")
            
        except Exception as e:
            logging.error(f"Failed to use local model: {str(e)}")
            # Return a helpful message about starting Ollama
            return ("I apologize, but I'm having trouble connecting to the local model. "
                   "Please ensure Ollama is running by opening a terminal and running 'ollama serve' "
                   "before asking questions.")
            
    except Exception as e:
        logging.error(f"Error in generate_llm_response: {str(e)}")
        return "I apologize, but I'm having trouble processing your question. Please try again."

@app.route('/')
def index():
    return render_template('chat.html')

@app.route('/upload', methods=['POST'])
@rate_limit()
def upload_file():
    try:
        if 'image' not in request.files:
            return jsonify({'error': 'No image provided'}), 400
        
        file = request.files['image']
        if file.filename == '':
            return jsonify({'error': 'No selected file'}), 400
        
        # Validate file
        try:
            validate_file(file)
        except ValueError as e:
            return jsonify({'error': str(e)}), 400
        
        if file:
            filename = secure_filename(file.filename)
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(filepath)
            
            # Generate a unique chat ID if not provided
            chat_id = request.form.get('chat_id')
            if not chat_id:
                chat_id = f"chat_{int(datetime.now().timestamp())}"
            
            try:
                # Start background processing
                asyncio.run(process_image_async(filepath, chat_id))
                
                # Encode the uploaded image to base64
                with open(filepath, "rb") as image_file:
                    encoded_string = base64.b64encode(image_file.read()).decode("utf-8")
                    base64_image = f"data:image/jpeg;base64,{encoded_string}"

                return jsonify({
                    'success': True,
                    'chat_id': chat_id,
                    'status': 'processing',
                    'image_data': base64_image
                })
                
            except Exception as e:
                logging.error(f"Error processing image: {str(e)}")
                return jsonify({'error': str(e)}), 500
    except Exception as e:
        logging.error(f"Error in upload_file: {str(e)}")
        return jsonify({'error': 'An error occurred while uploading the image'}), 500

@app.route('/ask', methods=['POST'])
@rate_limit()
def ask_question():
    try:
        data = request.json
        question = data.get('question')
        chat_id = data.get('chat_id')
        
        logging.info(f"Received question for chat_id: {chat_id}")
        
        if not question:
            return jsonify({'error': 'Missing question'}), 400
            
        if not chat_id:
            return jsonify({'error': 'Missing chat ID'}), 400
            
        # Get cached image data
        image_data = processed_image_data.get(chat_id)
        if not image_data:
            return jsonify({'error': 'No image data found for this chat. Please upload an image first.'}), 400
        
        try:
            # Generate response
            response = generate_llm_response(question, image_data)
            
            # Store Q&A in image_data
            if 'qa_history' not in image_data:
                image_data['qa_history'] = []
            
            image_data['qa_history'].append({
                'timestamp': datetime.now().isoformat(),
                'question': question,
                'answer': response
            })
            
            # Update timestamp and save
            chat_timestamps[chat_id] = datetime.now()
            save_chat_history()
            
            return jsonify({'answer': response})
        except Exception as e:
            logging.error(f"Error generating response: {str(e)}")
            return jsonify({'error': 'Error generating response. Please try again.'}), 500
            
    except Exception as e:
        logging.error(f"Error in ask_question: {str(e)}")
        return jsonify({'error': 'An error occurred while processing your question'}), 500

@app.route('/get_chat_history', methods=['GET'])
def get_chat_history():
    try:
        chat_id = request.args.get('chat_id')
        if chat_id and chat_id in processed_image_data:
            return jsonify({
                'success': True,
                'image_data': processed_image_data[chat_id]
            })
        return jsonify({'success': False, 'message': 'No chat history found'})
    except Exception as e:
        logging.error(f"Error in get_chat_history: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/get_all_chats', methods=['GET'])
def get_all_chats():
    try:
        # Convert timestamps to string format
        chat_data = {}
        for chat_id, data in processed_image_data.items():
            chat_data[chat_id] = {
                'timestamp': chat_timestamps[chat_id].isoformat() if chat_id in chat_timestamps else None,
                'qa_history': data.get('qa_history', []),
                'caption': data.get('caption', ''),
                'objects': data.get('objects', [])
            }
        
        return jsonify({
            'success': True,
            'chats': chat_data
        })
    except Exception as e:
        logging.error(f"Error in get_all_chats: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/delete_chat', methods=['POST'])
def delete_chat():
    """Delete a chat and its associated data"""
    try:
        chat_id = request.args.get('chat_id')
        if not chat_id:
            return jsonify({'error': 'Missing chat ID'}), 400
            
        # Remove from processed_image_data
        if chat_id in processed_image_data:
            del processed_image_data[chat_id]
            
        # Remove from chat_timestamps
        if chat_id in chat_timestamps:
            del chat_timestamps[chat_id]
            
        # Save changes
        save_chat_history()
        
        return jsonify({'success': True})
    except Exception as e:
        logging.error(f"Error in delete_chat: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/get_chat_image', methods=['GET'])
def get_chat_image():
    """Get the image for a specific chat"""
    try:
        chat_id = request.args.get('chat_id')
        if not chat_id or chat_id not in processed_image_data:
            return jsonify({'error': 'Chat not found'}), 404
            
        image_path = processed_image_data[chat_id].get('image_path')
        if not image_path or not os.path.exists(image_path):
            return jsonify({'error': 'Image not found'}), 404
            
        # Convert image to base64
        with open(image_path, "rb") as img_file:
            img_data = base64.b64encode(img_file.read()).decode('utf-8')
            return jsonify({
                'success': True,
                'image_data': f"data:image/{image_path.split('.')[-1]};base64,{img_data}"
            })
    except Exception as e:
        logging.error(f"Error in get_chat_image: {str(e)}")
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    # Ensure Ollama model is available
    try:
        ensure_ollama_model('deepseek-r1:latest')
    except Exception as e:
        logging.error(f"Failed to initialize Ollama model: {str(e)}")
        exit(1)
    
    # Start Ollama service only if not already running
    try:
        # Check if Ollama is running by attempting a simple API call
        ollama.list()
        logging.info("Ollama service is already running")
    except Exception:
        try:
            subprocess.run(['ollama', 'serve'], start_new_session=True)
            logging.info("Ollama service started successfully")
            # Wait briefly to ensure service is up
            time.sleep(2)
        except Exception as e:
            logging.warning(f"Could not start Ollama service: {str(e)}")
    
    app.run(debug=True)