Spaces:
Sleeping
Sleeping
| """ | |
| 3D Gaussian-style 4D Space Visualization | |
| ========================================= | |
| A web app that extracts evenly-spaced frames from an MP4 video and stacks them | |
| along the z-axis in a 3D scene, so the user can view the whole sequence of | |
| images at once as a "4D" volume of time. | |
| Backend responsibilities: | |
| * Accept an MP4 upload (or use the bundled sample video). | |
| * Extract N frames sampled evenly across the video's time domain (OpenCV). | |
| * Serve the extracted frames + a session manifest to the Three.js frontend. | |
| The frontend (static/index.html) does all the 3D rendering. | |
| """ | |
| import os | |
| import shutil | |
| import uuid | |
| import cv2 | |
| from flask import ( | |
| Flask, | |
| jsonify, | |
| request, | |
| send_from_directory, | |
| ) | |
| BASE_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| STATIC_DIR = os.path.join(BASE_DIR, "static") | |
| UPLOAD_DIR = os.path.join(BASE_DIR, "uploads") | |
| FRAMES_DIR = os.path.join(BASE_DIR, "frames") | |
| VIDEO_DIR = os.path.join(BASE_DIR, "video") | |
| def _find_sample_video(video_dir): | |
| """Return the path to the bundled sample video. | |
| Prefers common names (solar.mp4 / Lunar.mp4) but falls back to the first | |
| video file found in the folder, so renaming the sample won't break startup. | |
| """ | |
| preferred = ["solar.mp4", "Lunar.mp4", "sample.mp4"] | |
| for name in preferred: | |
| p = os.path.join(video_dir, name) | |
| if os.path.isfile(p): | |
| return p | |
| if os.path.isdir(video_dir): | |
| for name in sorted(os.listdir(video_dir)): | |
| if os.path.splitext(name)[1].lower() in {".mp4", ".mov", ".m4v", ".avi", ".mkv", ".webm"}: | |
| return os.path.join(video_dir, name) | |
| # Default path even if it doesn't exist yet (health check will report it). | |
| return os.path.join(video_dir, "solar.mp4") | |
| SAMPLE_VIDEO = _find_sample_video(VIDEO_DIR) | |
| ALLOWED_EXTENSIONS = {".mp4", ".mov", ".m4v", ".avi", ".mkv", ".webm"} | |
| MAX_CONTENT_LENGTH = 1024 * 1024 * 1024 # 1 GB upload cap | |
| def _ensure_writable(primary, fallback): | |
| """Return a writable directory, preferring `primary`, else `fallback`. | |
| On hosted environments (e.g. Hugging Face Spaces) the app directory may be | |
| read-only depending on the runtime user, so we fall back to a temp dir. | |
| """ | |
| for path in (primary, fallback): | |
| try: | |
| os.makedirs(path, exist_ok=True) | |
| testfile = os.path.join(path, ".write_test") | |
| with open(testfile, "w") as fh: | |
| fh.write("ok") | |
| os.remove(testfile) | |
| return path | |
| except OSError: | |
| continue | |
| return fallback | |
| import tempfile # noqa: E402 | |
| _TMP = tempfile.gettempdir() | |
| UPLOAD_DIR = _ensure_writable(UPLOAD_DIR, os.path.join(_TMP, "gauss4d_uploads")) | |
| FRAMES_DIR = _ensure_writable(FRAMES_DIR, os.path.join(_TMP, "gauss4d_frames")) | |
| app = Flask(__name__, static_folder=None) | |
| app.config["MAX_CONTENT_LENGTH"] = MAX_CONTENT_LENGTH | |
| # --------------------------------------------------------------------------- # | |
| # Frame extraction | |
| # --------------------------------------------------------------------------- # | |
| def extract_frames(video_path, session_id, num_frames, max_width=640): | |
| """Extract `num_frames` frames sampled evenly across the video timeline. | |
| Returns a manifest dict describing the extracted frames. Frames are written | |
| as JPEGs into frames/<session_id>/. | |
| """ | |
| cap = cv2.VideoCapture(video_path) | |
| if not cap.isOpened(): | |
| raise ValueError("Could not open video file. It may be corrupt or use an unsupported codec.") | |
| total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| fps = cap.get(cv2.CAP_PROP_FPS) or 0.0 | |
| src_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| src_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| # Some containers report a bogus/zero frame count. Fall back to a scan. | |
| if total <= 0: | |
| total = _count_frames_by_scan(cap) | |
| cap.release() | |
| cap = cv2.VideoCapture(video_path) | |
| if total <= 0: | |
| cap.release() | |
| raise ValueError("Video appears to contain no frames.") | |
| num_frames = max(1, min(int(num_frames), total)) | |
| out_dir = os.path.join(FRAMES_DIR, session_id) | |
| if os.path.isdir(out_dir): | |
| shutil.rmtree(out_dir) | |
| os.makedirs(out_dir, exist_ok=True) | |
| # Evenly spaced frame indices across the full timeline (inclusive of ends). | |
| if num_frames == 1: | |
| indices = [0] | |
| else: | |
| indices = [round(i * (total - 1) / (num_frames - 1)) for i in range(num_frames)] | |
| frames = [] | |
| out_w, out_h = src_w, src_h | |
| if src_w > max_width and src_w > 0: | |
| scale = max_width / float(src_w) | |
| out_w = max_width | |
| out_h = max(1, int(round(src_h * scale))) | |
| for i, frame_idx in enumerate(indices): | |
| cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx) | |
| ok, frame = cap.read() | |
| if not ok or frame is None: | |
| # Seeking failed at this index; skip gracefully. | |
| continue | |
| if (out_w, out_h) != (src_w, src_h): | |
| frame = cv2.resize(frame, (out_w, out_h), interpolation=cv2.INTER_AREA) | |
| fname = f"frame_{i:04d}.jpg" | |
| cv2.imwrite(os.path.join(out_dir, fname), frame, [cv2.IMWRITE_JPEG_QUALITY, 88]) | |
| frames.append( | |
| { | |
| "index": i, | |
| "sourceFrame": int(frame_idx), | |
| "time": (frame_idx / fps) if fps > 0 else None, | |
| "url": f"/frames/{session_id}/{fname}", | |
| } | |
| ) | |
| cap.release() | |
| if not frames: | |
| raise ValueError("Failed to extract any frames from the video.") | |
| aspect = (out_w / out_h) if out_h else 1.0 | |
| return { | |
| "sessionId": session_id, | |
| "count": len(frames), | |
| "requested": num_frames, | |
| "totalFrames": total, | |
| "fps": fps, | |
| "width": out_w, | |
| "height": out_h, | |
| "aspect": aspect, | |
| "frames": frames, | |
| } | |
| def _count_frames_by_scan(cap): | |
| count = 0 | |
| while True: | |
| ok = cap.grab() | |
| if not ok: | |
| break | |
| count += 1 | |
| return count | |
| # --------------------------------------------------------------------------- # | |
| # Routes | |
| # --------------------------------------------------------------------------- # | |
| def index(): | |
| return send_from_directory(STATIC_DIR, "index.html") | |
| def static_files(filename): | |
| return send_from_directory(STATIC_DIR, filename) | |
| def serve_frame(session_id, filename): | |
| return send_from_directory(os.path.join(FRAMES_DIR, session_id), filename) | |
| def api_extract(): | |
| """Extract frames from an uploaded video, or from the bundled sample. | |
| Form fields: | |
| * numFrames (int) - how many frames to sample (default 200) | |
| * useSample ("true"/"1") - use the bundled Lunar.mp4 instead of upload | |
| * video (file) - the uploaded video (when not using sample) | |
| """ | |
| try: | |
| num_frames = int(request.form.get("numFrames", 200)) | |
| except (TypeError, ValueError): | |
| num_frames = 200 | |
| num_frames = max(1, min(num_frames, 2000)) | |
| use_sample = str(request.form.get("useSample", "")).lower() in ("1", "true", "yes") | |
| session_id = uuid.uuid4().hex[:12] | |
| if use_sample: | |
| if not os.path.isfile(SAMPLE_VIDEO): | |
| return jsonify({"error": "Sample video not found on server."}), 404 | |
| video_path = SAMPLE_VIDEO | |
| else: | |
| file = request.files.get("video") | |
| if file is None or file.filename == "": | |
| return jsonify({"error": "No video file uploaded."}), 400 | |
| ext = os.path.splitext(file.filename)[1].lower() | |
| if ext not in ALLOWED_EXTENSIONS: | |
| return jsonify({"error": f"Unsupported file type: {ext}"}), 400 | |
| video_path = os.path.join(UPLOAD_DIR, f"{session_id}{ext}") | |
| file.save(video_path) | |
| try: | |
| manifest = extract_frames(video_path, session_id, num_frames) | |
| except ValueError as exc: | |
| return jsonify({"error": str(exc)}), 400 | |
| except Exception as exc: # pragma: no cover - defensive | |
| return jsonify({"error": f"Extraction failed: {exc}"}), 500 | |
| return jsonify(manifest) | |
| def health(): | |
| return jsonify({"status": "ok", "sampleAvailable": os.path.isfile(SAMPLE_VIDEO)}) | |
| if __name__ == "__main__": | |
| # Hugging Face Spaces expects the app on port 7860. Locally you can override | |
| # with the PORT env var. Bind to 0.0.0.0 so it's reachable inside a container. | |
| host = os.environ.get("HOST", "0.0.0.0") | |
| port = int(os.environ.get("PORT", "7860")) | |
| print("=" * 60) | |
| print(" 3D Gaussian 4D Video Stack Visualizer") | |
| print(f" Sample video: {SAMPLE_VIDEO}") | |
| print(f" Sample found: {os.path.isfile(SAMPLE_VIDEO)}") | |
| print(f" Serving on http://{host}:{port}") | |
| print(" (the sample will auto-extract and visualize on load)") | |
| print("=" * 60) | |
| app.run(host=host, port=port, debug=False, threaded=True) | |