Hướng dẫn xây dựng Hierarchical NeRF với JAX3D: Từ kết xuất thể tích đến tái tạo 3D

Trong bài hướng dẫn này, chúng ta sẽ xây dựng một Neural Radiance Field (NeRF) phân cấp từ đầu đến cuối (end-to-end) bằng cách sử dụng JAX, Flax, Optax và các nguyên hàm kết xuất thể tích (volume-rendering primitives) được cung cấp bởi jax3d. Trước tiên, chúng ta xây dựng một tập dữ liệu đa góc nhìn tổng hợp từ một cảnh phân tích chứa hình học thể tích và bức xạ phụ thuộc vào góc nhìn, sử dụng sample_along_rays và volume_rendering để thiết lập quy trình kết xuất thuận (forward rendering). Sau đó, chúng ta triển khai NeRF với mã hóa vị trí (positional encoding), kết nối tắt (skip connections), các mạng thô và tinh riêng biệt, cùng điều kiện hóa hướng nhìn (view-direction conditioning), theo sau là lấy mẫu quan trọng phân cấp (hierarchical importance sampling) thông qua sample_piecewise_constant_pdf. Chúng ta huấn luyện mô hình bằng JAX JIT compilation, tối ưu hóa Adam, phân rã tốc độ học theo hàm mũ (exponential learning-rate decay) và cắt gradient (gradient clipping), cuối cùng đánh giá khả năng tổng hợp góc nhìn mới (novel-view synthesis) bằng PSNR, trực quan hóa độ sâu và độ mờ, chẩn đoán lấy mẫu, kết xuất 360 độ và trích xuất hình học bằng marching-cubes.

Chúng ta thiết lập môi trường JAX3D, cài đặt các phụ thuộc cần thiết và tải module volume_rendering trực tiếp từ kho lưu trữ đã sao chép. Chúng ta cấu hình các tham số huấn luyện thích ứng GPU/CPU và thiết lập mô hình camera bằng các thông số nội tại pinhole (pinhole intrinsics), tư thế look-at và vị trí đặt camera dựa trên quỹ đạo. Sau đó, chúng ta tạo các tia không gian thế giới đã chuẩn hóa từ mỗi tư thế camera, cung cấp nền tảng hình học cho quy trình kết xuất.

Chúng ta xây dựng một cảnh phân tích thực tế (ground-truth scene) chứa các hình cầu cạnh mềm, sàn có họa tiết và bức xạ phản chiếu phụ thuộc vào góc nhìn. Chúng ta kết xuất cảnh này bằng trình triển khai volume-rendering của JAX3D để tạo ra các quan sát RGB, độ sâu và giá trị độ mờ nhất quán trên nhiều góc nhìn camera. Chúng ta tổ chức các hình ảnh thu được thành một tập hợp tia phẳng (flattened ray pool) để có thể lấy mẫu các tia ngẫu nhiên một cách hiệu quả trong quá trình huấn luyện NeRF.

Chúng ta triển khai biểu diễn NeRF bằng cách sử dụng mã hóa vị trí hình sin cho cả tọa độ không gian và hướng nhìn. Chúng ta sử dụng một MLP Flax sâu với kết nối tắt để dự đoán mật độ thể tích không âm từ vị trí, đồng thời điều kiện hóa RGB dựa trên hướng nhìn. Do đó, chúng ta tách biệt hình học độc lập với góc nhìn khỏi diện mạo phụ thuộc vào góc nhìn, cho phép mô hình biểu diễn cả cấu trúc cảnh và các hiệu ứng phản chiếu.

Chúng ta triển khai trình kết xuất phân cấp cốt lõi bằng cách lấy mẫu các điểm thô dọc theo mỗi tia và tổng hợp mật độ cũng như màu sắc của chúng thông qua toán tử volume-rendering của JAX3D. Chúng ta chuyển đổi trọng số kết xuất thô thu được thành phân phối xác suất hằng số từng đoạn và lấy mẫu quan trọng các điểm tinh bổ sung xung quanh các vùng có đóng góp cao. Chúng ta kết hợp và sắp xếp các mẫu thô và tinh trước khi thực hiện kết xuất mạng tinh cuối cùng, đồng thời dừng gradient thông qua thao tác lấy mẫu.

Chúng ta khởi tạo các mạng NeRF thô và tinh độc lập và tối ưu hóa chúng cùng nhau bằng Adam, sử dụng phân rã tốc độ học theo hàm mũ và cắt gradient toàn cục. Chúng ta giám sát cả hai giai đoạn kết xuất dựa trên màu sắc tia thực tế, khuyến khích mạng thô học các phân phối lấy mẫu hữu ích trong khi cải thiện quá trình tái tạo tinh cuối cùng. Chúng ta chạy bước huấn luyện với JAX JIT compilation và theo dõi PSNR của mạng tinh trong suốt quá trình tối ưu hóa.

Chúng ta đánh giá biểu diễn đã huấn luyện thông qua việc kết xuất góc nhìn mới theo từng phần (chunked) và đo lường chất lượng tái tạo bằng PSNR trên tập dữ liệu giữ lại, cùng với bản đồ độ sâu và độ mờ. Chúng ta trực quan hóa cách lấy mẫu phân cấp tập trung các mẫu tinh xung quanh các bề mặt quan trọng, sau đó tạo GIF quỹ đạo 360 độ để kiểm tra trường bức xạ đã học từ nhiều góc nhìn. Cuối cùng, chúng ta truy vấn mật độ đã học trên lưới 3D và áp dụng marching cubes để trích xuất một bề mặt đẳng trị hình học xấp xỉ.

Tóm lại, chúng ta đã chứng minh quy trình kết xuất ngược hoàn chỉnh bằng cách học một trường mật độ và bức xạ liên tục từ các quan sát đa góc nhìn tổng hợp và tái tạo nó thông qua kết xuất thể tích phân cấp. Chúng ta đã sử dụng mạng thô để xác định các vùng thông tin dọc theo mỗi tia và mạng tinh để tập trung các mẫu bổ sung xung quanh các bề mặt có đóng góp cao. Đồng thời, mã hóa hướng nhìn cho phép chúng ta mô hình hóa diện mạo phụ thuộc vào góc nhìn. Trong các giai đoạn đánh giá cuối cùng, chúng ta đã đo lường chất lượng tái tạo góc nhìn mới bằng PSNR, kiểm tra độ sâu và độ mờ đã học, trực quan hóa hành vi lấy mẫu quan trọng, tạo quỹ đạo 360 độ và trích xuất hình học xấp xỉ bằng marching cubes. Nhìn chung, chúng ta đã chỉ ra cách các thành phần toán học của jax3d tích hợp với quá trình huấn luyện mạng thần kinh dựa trên JAX hiện đại để tạo thành một hệ thống tái tạo NeRF nhỏ gọn nhưng hoàn thiện về mặt kỹ thuật.

Xem MÃ NGUỒN ĐẦY ĐỦ tại đây. Mọi tín dụng thuộc về nhà nghiên cứu của dự án này. Ngoài ra, hãy thoải mái theo dõi chúng tôi trên Twitter và đừng quên tham gia SubReddit ML 150k+ của chúng tôi và Đăng ký nhận Bản tin của chúng tôi. Khoan đã! bạn có dùng Telegram không? bây giờ bạn cũng có thể tham gia cùng chúng tôi trên Telegram.

Bạn cần hợp tác với chúng tôi để quảng bá GitHub Repo HOẶC Trang Hugging Face HOẶC Bản phát hành sản phẩm HOẶC Hội thảo trực tuyến, v.v.? Hãy kết nối với chúng tôi.