Spaces:
Runtime error
Runtime error
from flask import Flask, request, render_template, Response | |
import cv2 | |
import numpy as np | |
import tensorflow as tf | |
import threading | |
# Load the TFLite model | |
interpreter = tf.lite.Interpreter(model_path=r'midas.tflite') | |
interpreter.allocate_tensors() | |
input_details = interpreter.get_input_details() | |
output_details = interpreter.get_output_details() | |
app = Flask(__name__) | |
# Function to preprocess the image | |
def preprocess_image(image): | |
image = cv2.resize(image, (256, 256)) # Resize to 256x256 | |
image = image.astype(np.float32) / 255.0 # Normalize to [0,1] | |
image = np.expand_dims(image, axis=0) # Add batch dimension | |
return image | |
# Function to process the frame | |
def process_frame(frame): | |
input_image = preprocess_image(frame) | |
# Set the input tensor | |
interpreter.set_tensor(input_details[0]['index'], input_image) | |
# Run inference | |
interpreter.invoke() | |
# Get the output tensor | |
depth_map = interpreter.get_tensor(output_details[0]['index']) | |
# Process depth map | |
depth_map = np.squeeze(depth_map) | |
depth_map = (depth_map / np.max(depth_map) * 255).astype(np.uint8) | |
depth_map_gray = cv2.cvtColor(depth_map, cv2.COLOR_GRAY2BGR) | |
# Apply Canny edge detection on the original frame | |
edges = cv2.Canny(frame, threshold1=100, threshold2=200) | |
edges_colored = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR) | |
# Resize edges_colored to match depth_map dimensions | |
edges_colored = cv2.resize(edges_colored, depth_map_gray.shape[1::-1]) | |
# Create an overlay | |
alpha = 0.9 | |
overlay = cv2.addWeighted(depth_map_gray, alpha, edges_colored, alpha, 0) | |
overlay_gray = cv2.cvtColor(overlay, cv2.COLOR_BGR2GRAY) | |
# Segment processing | |
height, width = overlay_gray.shape | |
segment_width = width // 7 | |
avg_pixel_densities = [] | |
for i in range(7): | |
start_x = i * segment_width | |
end_x = (i + 1) * segment_width if i < 6 else width | |
segment_pixels = overlay_gray[:, start_x:end_x] | |
avg_pixel_density = np.mean(segment_pixels) | |
avg_pixel_densities.append(avg_pixel_density) | |
# Draw vertical lines and pixel density values | |
for i in range(7): | |
x = i * segment_width | |
cv2.line(overlay_gray, (x, 0), (x, height), (255, 255, 255), 1) | |
cv2.putText(overlay_gray, f"{avg_pixel_densities[i]:.2f}", (x + 5, 40), | |
cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1) | |
# Draw a dot and horizontal line | |
center_x = width // 2 | |
bottom_y = height - 10 | |
dot_position = (center_x, bottom_y) | |
cv2.circle(overlay_gray, dot_position, 5, (255, 255, 255), -1) | |
middle_y = height // 2 | |
cv2.line(overlay_gray, (0, middle_y), (width, middle_y), (255, 255, 255), 1) | |
# Draw arrow | |
lowest_index = np.argmin(avg_pixel_densities) | |
lowest_x = lowest_index * segment_width + segment_width // 2 | |
arrow_end = (lowest_x, height // 2) | |
cv2.arrowedLine(overlay_gray, dot_position, arrow_end, (255, 0, 0), 2, tipLength=0.1) | |
return overlay_gray | |
def index(): | |
return render_template('depthmap.html') | |
def video_feed(): | |
# Receive the frame from the client | |
frame_data = request.files['frame'].read() | |
nparr = np.frombuffer(frame_data, np.uint8) | |
frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR) | |
# Process the frame | |
processed_frame = process_frame(frame) | |
# Encode the processed frame as JPEG | |
_, jpeg = cv2.imencode('.jpg', processed_frame) | |
return Response(jpeg.tobytes(), mimetype='image/jpeg') | |
if __name__ == '__main__': | |
app.run(debug=True, threaded=True) | |