Spaces:
Build error
Build error
| """ | |
| Insta360 3D Reconstruction - Hugging Face Space Version | |
| Optimized for longer videos with intelligent frame sampling | |
| Supports ZeroGPU for faster processing | |
| """ | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from transformers import DPTForDepthEstimation, DPTImageProcessor | |
| import open3d as o3d | |
| import plotly.graph_objects as go | |
| import cv2 | |
| import tempfile | |
| from pathlib import Path | |
| import time | |
| import warnings | |
| from scipy import ndimage | |
| from scipy.ndimage import gaussian_filter | |
| import spaces # For ZeroGPU support | |
| warnings.filterwarnings('ignore') | |
| # Load model | |
| print("🔄 Loading depth estimation model...") | |
| try: | |
| dpt_processor = DPTImageProcessor.from_pretrained("Intel/dpt-large") | |
| dpt_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large") | |
| # Don't move to CUDA here - ZeroGPU will handle it in decorated functions | |
| dpt_model.eval() | |
| print("✅ Model loaded successfully! (ZeroGPU will handle GPU allocation)") | |
| except Exception as e: | |
| print(f"❌ Error loading model: {e}") | |
| dpt_processor = None | |
| dpt_model = None | |
| # Enhanced depth processing functions | |
| def bilateral_filter_depth(depth_map, d=9, sigma_color=75, sigma_space=75): | |
| """Apply bilateral filter to preserve edges while smoothing depth""" | |
| depth_norm = ((depth_map - depth_map.min()) / (depth_map.max() - depth_map.min()) * 255).astype(np.uint8) | |
| filtered = cv2.bilateralFilter(depth_norm, d, sigma_color, sigma_space) | |
| filtered = filtered.astype(np.float32) / 255.0 | |
| filtered = filtered * (depth_map.max() - depth_map.min()) + depth_map.min() | |
| return filtered | |
| def multi_scale_depth_refinement(depth_map, scales=[1.0, 0.5]): | |
| """Process depth at multiple scales and fuse""" | |
| h, w = depth_map.shape | |
| refined_depths = [] | |
| weights = [] | |
| for scale in scales: | |
| if scale == 1.0: | |
| scaled_depth = depth_map | |
| else: | |
| new_h, new_w = int(h * scale), int(w * scale) | |
| scaled_depth = cv2.resize(depth_map, (new_w, new_h), interpolation=cv2.INTER_LINEAR) | |
| scaled_depth = cv2.resize(scaled_depth, (w, h), interpolation=cv2.INTER_LINEAR) | |
| filtered_depth = bilateral_filter_depth(scaled_depth) | |
| refined_depths.append(filtered_depth) | |
| weights.append(scale) | |
| weights = np.array(weights) | |
| weights = weights / weights.sum() | |
| final_depth = np.zeros_like(depth_map) | |
| for depth, weight in zip(refined_depths, weights): | |
| final_depth += depth * weight | |
| return final_depth | |
| def estimate_depth_confidence(depth_map): | |
| """Estimate confidence map based on depth consistency""" | |
| grad_x = cv2.Sobel(depth_map, cv2.CV_64F, 1, 0, ksize=3) | |
| grad_y = cv2.Sobel(depth_map, cv2.CV_64F, 0, 1, ksize=3) | |
| grad_mag = np.sqrt(grad_x**2 + grad_y**2) | |
| confidence = 1.0 / (1.0 + grad_mag / grad_mag.max()) | |
| confidence = gaussian_filter(confidence, sigma=2) | |
| return confidence | |
| def intelligent_frame_sampling(video_path, target_frames=6, max_frames=100): | |
| """ | |
| Intelligently sample frames from video based on motion and content | |
| For long videos, this prevents processing too many similar frames | |
| """ | |
| cap = cv2.VideoCapture(video_path) | |
| total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| duration = total_frames / fps if fps > 0 else 0 | |
| # For very long videos, sample more intelligently | |
| if duration > 120: # 2 minutes | |
| # Sample every N seconds instead of uniformly | |
| sample_interval = max(int(fps * 15), 1) # Every 15 seconds | |
| frame_indices = list(range(0, total_frames, sample_interval)) | |
| else: | |
| # Uniform sampling | |
| frame_indices = np.linspace(0, total_frames - 1, min(target_frames, total_frames), dtype=int) | |
| # Limit to max_frames to prevent timeout | |
| if len(frame_indices) > max_frames: | |
| frame_indices = frame_indices[::len(frame_indices)//max_frames][:max_frames] | |
| cap.release() | |
| return frame_indices, total_frames, fps, duration | |
| def extract_frames_smart(video_path, target_frames=6): | |
| """Extract frames intelligently based on video length""" | |
| frame_indices, total_frames, fps, duration = intelligent_frame_sampling(video_path, target_frames) | |
| cap = cv2.VideoCapture(video_path) | |
| frames = [] | |
| for idx in frame_indices: | |
| cap.set(cv2.CAP_PROP_POS_FRAMES, idx) | |
| ret, frame = cap.read() | |
| if ret: | |
| frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| frames.append(frame_rgb) | |
| cap.release() | |
| info = { | |
| 'total_frames': total_frames, | |
| 'extracted_frames': len(frames), | |
| 'fps': fps, | |
| 'duration': duration, | |
| 'frame_indices': frame_indices | |
| } | |
| return frames, info | |
| def equirectangular_to_perspective(equirect_img, fov=90, theta=0, phi=0, height=512, width=512): | |
| """Convert equirectangular image to perspective view""" | |
| equ_h, equ_w = equirect_img.shape[:2] | |
| y, x = np.meshgrid(np.arange(height), np.arange(width), indexing='ij') | |
| x_norm = (2.0 * x / width - 1.0) | |
| y_norm = (2.0 * y / height - 1.0) | |
| fov_rad = np.radians(fov) | |
| focal = 0.5 * width / np.tan(0.5 * fov_rad) | |
| z_cam = focal | |
| x_cam = x_norm * width / 2.0 | |
| y_cam = y_norm * height / 2.0 | |
| norm = np.sqrt(x_cam**2 + y_cam**2 + z_cam**2) | |
| x_cam /= norm | |
| y_cam /= norm | |
| z_cam /= norm | |
| theta_rad = np.radians(theta) | |
| phi_rad = np.radians(phi) | |
| rot_y = np.array([ | |
| [np.cos(theta_rad), 0, np.sin(theta_rad)], | |
| [0, 1, 0], | |
| [-np.sin(theta_rad), 0, np.cos(theta_rad)] | |
| ]) | |
| rot_x = np.array([ | |
| [1, 0, 0], | |
| [0, np.cos(phi_rad), -np.sin(phi_rad)], | |
| [0, np.sin(phi_rad), np.cos(phi_rad)] | |
| ]) | |
| rot = rot_y @ rot_x | |
| rays = np.stack([x_cam, y_cam, z_cam], axis=-1) | |
| rays_rot = rays @ rot.T | |
| x_rot = rays_rot[..., 0] | |
| y_rot = rays_rot[..., 1] | |
| z_rot = rays_rot[..., 2] | |
| lon = np.arctan2(x_rot, z_rot) | |
| lat = np.arcsin(np.clip(y_rot, -1, 1)) | |
| equ_x = (lon / np.pi + 1) * 0.5 * (equ_w - 1) | |
| equ_y = (0.5 - lat / np.pi) * (equ_h - 1) | |
| equ_x = np.clip(equ_x, 0, equ_w - 1) | |
| equ_y = np.clip(equ_y, 0, equ_h - 1) | |
| perspective_img = np.zeros((height, width, equirect_img.shape[2]), dtype=equirect_img.dtype) | |
| for c in range(equirect_img.shape[2]): | |
| perspective_img[..., c] = ndimage.map_coordinates( | |
| equirect_img[..., c], | |
| [equ_y, equ_x], | |
| order=1, | |
| mode='wrap' | |
| ) | |
| return perspective_img | |
| # ZeroGPU decorator for GPU acceleration | |
| def estimate_depth_enhanced(image, processor, model): | |
| """Enhanced depth estimation with multi-scale processing""" | |
| inputs = processor(images=image, return_tensors="pt") | |
| # ZeroGPU automatically handles device placement | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| predicted_depth = outputs.predicted_depth | |
| prediction = torch.nn.functional.interpolate( | |
| predicted_depth.unsqueeze(1), | |
| size=image.shape[:2], | |
| mode="bicubic", | |
| align_corners=False, | |
| ) | |
| depth_map = prediction.squeeze().cpu().numpy() | |
| depth_map = multi_scale_depth_refinement(depth_map) | |
| confidence = estimate_depth_confidence(depth_map) | |
| return depth_map, confidence | |
| def depth_to_point_cloud_enhanced(depth, color, confidence, camera_params): | |
| """Enhanced point cloud generation with confidence weighting""" | |
| height, width = depth.shape | |
| fx, fy = camera_params['fx'], camera_params['fy'] | |
| cx, cy = camera_params['cx'], camera_params['cy'] | |
| R_matrix = camera_params.get('R', np.eye(3)) | |
| t_vector = camera_params.get('t', np.zeros(3)) | |
| u, v = np.meshgrid(np.arange(width), np.arange(height)) | |
| z = depth | |
| x = (u - cx) * z / fx | |
| y = (v - cy) * z / fy | |
| points_cam = np.stack([x, y, z], axis=-1) | |
| points_world = points_cam @ R_matrix.T + t_vector | |
| conf_threshold = np.percentile(confidence, 30) | |
| valid_mask = confidence > conf_threshold | |
| points = points_world[valid_mask] | |
| colors = color[valid_mask] | |
| return points, colors | |
| def create_realistic_mesh(points, colors, progress_callback): | |
| """Create high-quality mesh using Poisson reconstruction""" | |
| progress_callback("🎨 Creating realistic mesh...") | |
| pcd = o3d.geometry.PointCloud() | |
| pcd.points = o3d.utility.Vector3dVector(points) | |
| pcd.colors = o3d.utility.Vector3dVector(colors / 255.0) | |
| progress_callback(" • Removing outliers...") | |
| pcd, _ = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0) | |
| progress_callback(" • Estimating normals...") | |
| pcd.estimate_normals( | |
| search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30) | |
| ) | |
| pcd.orient_normals_consistent_tangent_plane(k=15) | |
| progress_callback(" • Performing Poisson reconstruction...") | |
| mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson( | |
| pcd, depth=9, width=0, scale=1.1, linear_fit=False | |
| ) | |
| progress_callback(" • Cleaning mesh...") | |
| densities = np.asarray(densities) | |
| density_threshold = np.percentile(densities, 10) | |
| vertices_to_remove = densities < density_threshold | |
| mesh.remove_vertices_by_mask(vertices_to_remove) | |
| mesh = mesh.filter_smooth_simple(number_of_iterations=5) | |
| mesh.compute_vertex_normals() | |
| # Transfer colors | |
| mesh_points = np.asarray(mesh.vertices) | |
| pcd_tree = o3d.geometry.KDTreeFlann(pcd) | |
| pcd_colors = np.asarray(pcd.colors) | |
| mesh_colors = np.zeros_like(mesh_points) | |
| for i, point in enumerate(mesh_points): | |
| [_, idx, _] = pcd_tree.search_knn_vector_3d(point, 1) | |
| mesh_colors[i] = pcd_colors[idx[0]] | |
| mesh.vertex_colors = o3d.utility.Vector3dVector(mesh_colors) | |
| return mesh | |
| def process_video(video_path, num_frames, num_views, quality, progress=gr.Progress()): | |
| """Main processing function optimized for Hugging Face""" | |
| if dpt_model is None: | |
| return None, None, None, "❌ Model not loaded properly", None | |
| if video_path is None: | |
| return None, None, None, "❌ Please upload a video first", None | |
| status = [] | |
| start_time = time.time() | |
| def update_status(msg): | |
| status.append(msg) | |
| progress(0.1, desc=msg) | |
| return "\n".join(status) | |
| try: | |
| status_text = update_status("="*60) | |
| status_text = update_status("🎬 STARTING REALISTIC 3D RECONSTRUCTION") | |
| status_text = update_status("="*60) | |
| # Check video | |
| cap = cv2.VideoCapture(video_path) | |
| if not cap.isOpened(): | |
| return None, None, None, "❌ Cannot open video file", None | |
| total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| duration = total_frames / fps if fps > 0 else 0 | |
| cap.release() | |
| status_text = update_status(f"\n📹 Video Info:") | |
| status_text = update_status(f" • Duration: {duration:.1f}s ({total_frames} frames)") | |
| status_text = update_status(f" • FPS: {fps:.1f}") | |
| # Warn about long videos | |
| if duration > 300: | |
| status_text = update_status(f"\n⚠️ WARNING: Very long video ({duration:.0f}s)") | |
| status_text = update_status(f" • Processing will be slower") | |
| status_text = update_status(f" • Consider using a shorter clip") | |
| # Extract frames intelligently | |
| status_text = update_status(f"\n📹 Extracting frames intelligently...") | |
| frames, video_info = extract_frames_smart(video_path, num_frames) | |
| if not frames: | |
| return None, None, None, "❌ Failed to extract frames", None | |
| status_text = update_status(f"✅ Extracted {len(frames)} frames") | |
| status_text = update_status(f" • Sampling strategy: {'Intelligent (long video)' if duration > 120 else 'Uniform'}") | |
| preview_img = Image.fromarray(frames[0]) | |
| # Quality settings | |
| quality_configs = { | |
| 'low': {'resolution': 384, 'fov': 90}, | |
| 'medium': {'resolution': 512, 'fov': 90}, | |
| 'high': {'resolution': 640, 'fov': 85} | |
| } | |
| config = quality_configs[quality] | |
| status_text = update_status(f"\n⚙️ Settings: {len(frames)} frames × {num_views} views × {config['resolution']}px") | |
| # Process frames | |
| all_points = [] | |
| all_colors = [] | |
| total_views = len(frames) * num_views | |
| processed_views = 0 | |
| for frame_idx, frame in enumerate(frames): | |
| progress((frame_idx + 1) / len(frames), desc=f"Processing frame {frame_idx+1}/{len(frames)}") | |
| status_text = update_status(f"\n📐 Frame {frame_idx + 1}/{len(frames)}:") | |
| # Generate view angles | |
| view_angles = [(360.0 / num_views * i, 0) for i in range(num_views)] | |
| frame_points = [] | |
| frame_colors = [] | |
| for view_idx, (theta, phi) in enumerate(view_angles): | |
| # Convert to perspective | |
| persp_img = equirectangular_to_perspective( | |
| frame, fov=config['fov'], theta=theta, phi=phi, | |
| height=config['resolution'], width=config['resolution'] | |
| ) | |
| # Depth estimation | |
| depth_map, confidence = estimate_depth_enhanced(persp_img, dpt_processor, dpt_model) | |
| # Camera params | |
| focal = config['resolution'] / (2 * np.tan(np.radians(config['fov']) / 2)) | |
| from scipy.spatial.transform import Rotation as R | |
| rot = R.from_euler('yz', [theta, phi], degrees=True) | |
| R_matrix = rot.as_matrix() | |
| camera_params = { | |
| 'fx': focal, 'fy': focal, | |
| 'cx': config['resolution'] / 2, | |
| 'cy': config['resolution'] / 2, | |
| 'R': R_matrix, | |
| 't': np.zeros(3) | |
| } | |
| # Generate points | |
| points, colors = depth_to_point_cloud_enhanced( | |
| depth_map, persp_img, confidence, camera_params | |
| ) | |
| frame_points.append(points) | |
| frame_colors.append(colors) | |
| processed_views += 1 | |
| if (view_idx + 1) % 2 == 0: | |
| status_text = update_status(f" • Processed {view_idx + 1}/{num_views} views") | |
| all_points.append(np.vstack(frame_points)) | |
| all_colors.append(np.vstack(frame_colors)) | |
| # Combine all | |
| status_text = update_status(f"\n🔗 Combining {len(frames)} frames...") | |
| final_points = np.vstack(all_points) | |
| final_colors = np.vstack(all_colors) | |
| status_text = update_status(f"✅ Total points: {len(final_points):,}") | |
| # Filter | |
| status_text = update_status(f"\n🎯 Filtering and cleaning...") | |
| # Remove duplicates | |
| unique_indices = np.unique(final_points, axis=0, return_index=True)[1] | |
| final_points = final_points[unique_indices] | |
| final_colors = final_colors[unique_indices] | |
| # Statistical outlier removal | |
| pcd_temp = o3d.geometry.PointCloud() | |
| pcd_temp.points = o3d.utility.Vector3dVector(final_points) | |
| pcd_temp, inlier_indices = pcd_temp.remove_statistical_outlier(nb_neighbors=30, std_ratio=2.0) | |
| final_points = final_points[inlier_indices] | |
| final_colors = final_colors[inlier_indices] | |
| status_text = update_status(f"✅ Filtered to {len(final_points):,} points") | |
| # Downsample if huge | |
| if len(final_points) > 500000: | |
| keep_ratio = 500000 / len(final_points) | |
| keep_indices = np.random.choice(len(final_points), size=int(len(final_points) * keep_ratio), replace=False) | |
| final_points = final_points[keep_indices] | |
| final_colors = final_colors[keep_indices] | |
| status_text = update_status(f" • Downsampled to {len(final_points):,} points") | |
| # Visualization | |
| status_text = update_status(f"\n📊 Creating 3D visualization...") | |
| vis_sample = min(50000, len(final_points)) | |
| vis_indices = np.random.choice(len(final_points), vis_sample, replace=False) | |
| vis_points = final_points[vis_indices] | |
| vis_colors = final_colors[vis_indices] | |
| fig = go.Figure(data=[go.Scatter3d( | |
| x=vis_points[:, 0], y=vis_points[:, 1], z=vis_points[:, 2], | |
| mode='markers', | |
| marker=dict( | |
| size=2, | |
| color=[f'rgb({int(c[0])},{int(c[1])},{int(c[2])})' for c in vis_colors], | |
| opacity=0.8 | |
| ) | |
| )]) | |
| fig.update_layout( | |
| title=f"3D Reconstruction ({len(final_points):,} points)", | |
| scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z', aspectmode='data'), | |
| height=700 | |
| ) | |
| # Save point cloud | |
| status_text = update_status(f"\n💾 Saving outputs...") | |
| pcd = o3d.geometry.PointCloud() | |
| pcd.points = o3d.utility.Vector3dVector(final_points) | |
| pcd.colors = o3d.utility.Vector3dVector(final_colors / 255.0) | |
| pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30)) | |
| ply_path = Path(tempfile.mkdtemp()) / "reconstruction.ply" | |
| o3d.io.write_point_cloud(str(ply_path), pcd) | |
| ply_path = str(ply_path) # Convert Path to string for Gradio | |
| status_text = update_status(f"✅ Point cloud saved") | |
| # Create mesh | |
| obj_path = None | |
| elapsed = time.time() - start_time | |
| if elapsed < 180: # Only if under 3 minutes so far | |
| try: | |
| def mesh_progress(msg): | |
| nonlocal status_text | |
| status_text = update_status(msg) | |
| mesh = create_realistic_mesh(final_points, final_colors, mesh_progress) | |
| obj_path = Path(tempfile.mkdtemp()) / "reconstruction.obj" | |
| o3d.io.write_triangle_mesh(str(obj_path), mesh) | |
| obj_path = str(obj_path) # Convert Path to string for Gradio | |
| status_text = update_status(f"✅ Mesh created: {len(mesh.vertices):,} vertices") | |
| except Exception as e: | |
| status_text = update_status(f"⚠️ Mesh generation failed: {str(e)}") | |
| else: | |
| status_text = update_status("⚠️ Mesh skipped (time limit)") | |
| # Final stats | |
| elapsed = time.time() - start_time | |
| status_text = update_status(f"\n{'='*60}") | |
| status_text = update_status(f"🎉 SUCCESS! Completed in {elapsed:.1f}s") | |
| status_text = update_status(f"📊 Final: {len(final_points):,} points") | |
| status_text = update_status(f"{'='*60}") | |
| return fig, ply_path, obj_path, status_text, preview_img | |
| except Exception as e: | |
| import traceback | |
| error_msg = f"❌ ERROR: {str(e)}\n\n{traceback.format_exc()}" | |
| return None, None, None, error_msg, None | |
| # Create Gradio interface | |
| with gr.Blocks(title="Insta360 3D Reconstruction", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(""" | |
| # 🌍 Insta360 3D Reconstruction | |
| ### Transform 360° videos into realistic 3D models | |
| **Optimized for videos of any length** - Uses intelligent frame sampling for longer videos | |
| """) | |
| gr.Markdown(""" | |
| ### ⚠️ For 8-Minute Videos: | |
| - Processing will take 10-15 minutes | |
| - Uses intelligent frame sampling (every 15 seconds) | |
| - Recommended: Use lower quality settings first | |
| - Consider trimming to 1-2 minutes for faster results | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| video_input = gr.Video(label="Upload 360° Video") | |
| with gr.Accordion("⚙️ Settings", open=True): | |
| num_frames = gr.Slider( | |
| minimum=4, maximum=12, value=6, step=2, | |
| label="Target Frames (auto-adjusted for long videos)" | |
| ) | |
| num_views = gr.Slider( | |
| minimum=4, maximum=8, value=6, step=2, | |
| label="Views per Frame" | |
| ) | |
| quality = gr.Radio( | |
| choices=['low', 'medium', 'high'], | |
| value='medium', | |
| label="Quality (Start with 'medium' for 8-min videos)" | |
| ) | |
| process_btn = gr.Button("🚀 Start Reconstruction", variant="primary", size="lg") | |
| with gr.Column(): | |
| status_output = gr.Textbox(label="Processing Status", lines=20, max_lines=25) | |
| preview_output = gr.Image(label="Video Preview") | |
| with gr.Row(): | |
| visualization_output = gr.Plot(label="3D Visualization") | |
| with gr.Row(): | |
| ply_output = gr.File(label="📦 Download Point Cloud (.ply)") | |
| obj_output = gr.File(label="📦 Download Mesh (.obj)") | |
| process_btn.click( | |
| fn=process_video, | |
| inputs=[video_input, num_frames, num_views, quality], | |
| outputs=[visualization_output, ply_output, obj_output, status_output, preview_output] | |
| ) | |
| gr.Markdown(""" | |
| ### 💡 Tips for Best Results | |
| **For 8-minute videos:** | |
| - Start with Medium quality (faster) | |
| - Uses intelligent sampling (~ every 15 seconds) | |
| - Total processing: 10-15 minutes | |
| - Or trim to 1-2 minutes for 3-5 min processing | |
| **Quality Guide:** | |
| - **Low**: 2-4 min (quick preview) | |
| - **Medium**: 5-10 min (good balance) | |
| - **High**: 10-20 min (best quality) | |
| **Video Requirements:** | |
| - Format: MP4 (equirectangular 360°) | |
| - Aspect Ratio: 2:1 | |
| - Any length (optimized for long videos) | |
| """) | |
| if __name__ == "__main__": | |
| demo.launch() |