Spaces:
Sleeping
Sleeping
File size: 2,300 Bytes
d9f749e c7c40ce d9f749e d8cfbd7 c7c40ce d9f749e c7c40ce |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import cv2
import torch
import numpy as np
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForDepthEstimation
import streamlit as st
# Load model and image processor
image_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
# Set the device for model (CUDA if available)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
# Use FP16 if available (half precision for speed)
if torch.cuda.is_available():
model = model.half()
# Streamlit App
st.title("Real-time Depth Estimation from Webcam")
# Initialize the webcam capture (OpenCV)
cap = cv2.VideoCapture(0)
# Streamlit button to capture a screenshot
if st.button("Capture Screenshot"):
ret, frame = cap.read()
if ret:
# Process the frame for depth estimation
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image = Image.fromarray(frame_rgb)
# Prepare image for the model
inputs = image_processor(images=image, return_tensors="pt").to(device)
# Model inference (no gradients needed)
with torch.no_grad():
outputs = model(**inputs)
predicted_depth = outputs.predicted_depth
# Interpolate depth map to match the frame's dimensions
prediction = torch.nn.functional.interpolate(
predicted_depth.unsqueeze(1),
size=(frame.shape[0], frame.shape[1]), # Match the frame's dimensions
mode="bicubic",
align_corners=False,
)
# Convert depth map to numpy for visualization
depth_map = prediction.squeeze().cpu().numpy()
# Normalize depth map for display (visualization purposes)
depth_map_normalized = np.uint8(depth_map / np.max(depth_map) * 255)
depth_map_colored = cv2.applyColorMap(depth_map_normalized, cv2.COLORMAP_JET)
# Display the original frame and the depth map in Streamlit
st.image(frame, caption="Original Webcam Image", channels="BGR", use_column_width=True)
st.image(depth_map_colored, caption="Depth Map", channels="BGR", use_column_width=True)
# Release the capture object when done
cap.release()
|