from flask import Flask, request, render_template, jsonify
import cv2
import numpy as np
import torch
from torchvision import transforms
import base64
from io import BytesIO
from PIL import Image
import threading
import queue

# Load the MiDaS model from PyTorch Hub
model = torch.hub.load("intel-isl/MiDaS", "MiDaS_small", trust_repo=True)
model.eval()

# Image transformation function
transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Create Flask app
app = Flask(__name__)

# Function to estimate depth from a frame and apply color mapping
def estimate_depth(frame):
    input_batch = transform(frame).unsqueeze(0)
    with torch.no_grad():
        prediction = model(input_batch)
        depth_map = prediction.squeeze().cpu().numpy()
        
        # Normalize and apply a colormap
        depth_map = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX)
        depth_map = depth_map.astype(np.uint8)
        colored_depth_map = cv2.applyColorMap(depth_map, cv2.COLORMAP_JET)
        return colored_depth_map

# Function to process the video frame in a separate thread
def process_frame_thread(data, response_queue):
    image_data = base64.b64decode(data.split(',')[1])
    image = Image.open(BytesIO(image_data))
    frame = np.array(image)

    # Convert RGB to BGR format (as OpenCV expects BGR)
    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    depth_map = estimate_depth(frame)

    # Encode depth map as a base64 image to send back
    _, buffer = cv2.imencode('.jpg', depth_map)
    depth_map_base64 = base64.b64encode(buffer).decode('utf-8')

    # Add the result to the response queue
    response_queue.put(f"data:image/jpeg;base64,{depth_map_base64}")

# Route to serve the HTML template
@app.route('/')
def index():
    return render_template('index.html')

# Route to process video frames and return depth map
@app.route('/process_frame', methods=['POST'])
def process_frame():
    data = request.json['image']
    
    # Create a queue to hold the response from the background thread
    response_queue = queue.Queue()
    
    # Start the processing thread
    thread = threading.Thread(target=process_frame_thread, args=(data, response_queue))
    thread.start()
    
    # Wait for the thread to complete and get the result from the queue
    thread.join()
    depth_map_base64 = response_queue.get()

    return jsonify({'depth_map': depth_map_base64})

if __name__ == "__main__":
    app.run(debug=True)