【三维重建探秘】NeRF：用神经网络捕捉光与影的艺术

import java.util.*;
import java.util.concurrent.*;
import java.awt.image.*;
import java.io.*;
import javax.imageio.*;

// 三维向量类
class Vector3 {
    double x, y, z;

    public Vector3(double x, double y, double z) {
        this.x = x;
        this.y = y;
        this.z = z;
    }

    public Vector3 add(Vector3 v) {
        return new Vector3(x + v.x, y + v.y, z + v.z);
    }

    public Vector3 subtract(Vector3 v) {
        return new Vector3(x - v.x, y - v.y, z - v.z);
    }

    public Vector3 multiply(double scalar) {
        return new Vector3(x * scalar, y * scalar, z * scalar);
    }

    public double dot(Vector3 v) {
        return x * v.x + y * v.y + z * v.z;
    }

    public Vector3 normalize() {
        double length = Math.sqrt(x * x + y * y + z * z);
        return new Vector3(x / length, y / length, z / length);
    }

    @Override
    public String toString() {
        return String.format("(%.4f, %.4f, %.4f)", x, y, z);
    }
}

// 颜色类
class Color {
    double r, g, b;

    public Color(double r, double g, double b) {
        this.r = r;
        this.g = g;
        this.b = b;
    }

    public Color add(Color c) {
        return new Color(r + c.r, g + c.g, b + c.b);
    }

    public Color multiply(double scalar) {
        return new Color(r * scalar, g * scalar, b * scalar);
    }

    public Color blend(Color c, double alpha) {
        return this.multiply(1 - alpha).add(c.multiply(alpha));
    }

    public int toRGB() {
        int rInt = (int) (Math.min(1.0, r) * 255);
        int gInt = (int) (Math.min(1.0, g) * 255);
        int bInt = (int) (Math.min(1.0, b) * 255);
        return (rInt << 16) | (gInt << 8) | bInt;
    }
}

// 神经网络层
class Layer {
    double[][] weights;
    double[] biases;
    String activation;

    public Layer(int inputSize, int outputSize, String activation) {
        this.weights = new double[inputSize][outputSize];
        this.biases = new double[outputSize];
        this.activation = activation;

        // 随机初始化权重和偏置
        Random random = new Random();
        for (int i = 0; i < inputSize; i++) {
            for (int j = 0; j < outputSize; j++) {
                weights[i][j] = random.nextGaussian() * 0.01;
            }
        }
        for (int j = 0; j < outputSize; j++) {
            biases[j] = random.nextGaussian() * 0.01;
        }
    }

    public double[] forward(double[] input) {
        double[] output = new double[weights[0].length];
        
        // 线性变换
        for (int j = 0; j < weights[0].length; j++) {
            double sum = biases[j];
            for (int i = 0; i < input.length; i++) {
                sum += input[i] * weights[i][j];
            }
            output[j] = sum;
        }
        
        // 激活函数
        if (activation.equals("relu")) {
            for (int j = 0; j < output.length; j++) {
                output[j] = Math.max(0, output[j]);
            }
        } else if (activation.equals("sigmoid")) {
            for (int j = 0; j < output.length; j++) {
                output[j] = 1.0 / (1.0 + Math.exp(-output[j]));
            }
        }
        
        return output;
    }
}

// 神经网络
class MLP {
    List<Layer> layers;

    public MLP() {
        this.layers = new ArrayList<>();
    }

    public void addLayer(Layer layer) {
        layers.add(layer);
    }

    public double[] forward(double[] input) {
        double[] output = input;
        for (Layer layer : layers) {
            output = layer.forward(output);
        }
        return output;
    }
}

// 射线类
class Ray {
    Vector3 origin;
    Vector3 direction;

    public Ray(Vector3 origin, Vector3 direction) {
        this.origin = origin;
        this.direction = direction.normalize();
    }

    public Vector3 pointAt(double t) {
        return origin.add(direction.multiply(t));
    }
}

// NeRF模型
class NeRF {
    MLP network;
    int numSamples;
    double near;
    double far;

    public NeRF(int numSamples, double near, double far) {
        this.numSamples = numSamples;
        this.near = near;
        this.far = far;
        
        // 构建神经网络
        network = new MLP();
        network.addLayer(new Layer(9, 256, "relu")); // 3D位置(3) + 方向(3) + 位置编码(3)
        network.addLayer(new Layer(256, 256, "relu"));
        network.addLayer(new Layer(256, 256, "relu"));
        network.addLayer(new Layer(256, 4, "sigmoid")); // 输出: RGB(3) + 密度(1)
    }

    // 位置编码函数
    public double[] positionalEncoding(double x, int L) {
        double[] encoded = new double[2 * L];
        for (int i = 0; i < L; i++) {
            encoded[2 * i] = Math.sin(Math.pow(2, i) * x);
            encoded[2 * i + 1] = Math.cos(Math.pow(2, i) * x);
        }
        return encoded;
    }

    // 渲染射线
    public Color renderRay(Ray ray) {
        // 在射线上均匀采样
        double[] tValues = new double[numSamples];
        for (int i = 0; i < numSamples; i++) {
            tValues[i] = near + (far - near) * (i + 0.5) / numSamples;
        }
        
        Color color = new Color(0, 0, 0);
        double transmittance = 1.0;
        
        // 对每个采样点进行评估
        for (int i = 0; i < numSamples; i++) {
            Vector3 point = ray.pointAt(tValues[i]);
            
            // 位置编码
            double[] encodedPos = new double[9];
            encodedPos[0] = point.x;
            encodedPos[1] = point.y;
            encodedPos[2] = point.z;
            
            double[] encodedX = positionalEncoding(point.x, 1);
            double[] encodedY = positionalEncoding(point.y, 1);
            double[] encodedZ = positionalEncoding(point.z, 1);
            
            System.arraycopy(encodedX, 0, encodedPos, 3, 2);
            System.arraycopy(encodedY, 0, encodedPos, 5, 2);
            System.arraycopy(encodedZ, 0, encodedPos, 7, 2);
            
            // 添加观察方向
            double[] input = new double[9];
            System.arraycopy(encodedPos, 0, input, 0, 9);
            
            // 前向传播
            double[] output = network.forward(input);
            
            // 提取颜色和密度
            Color sampleColor = new Color(output[0], output[1], output[2]);
            double density = output[3];
            
            // 计算透明度
            double delta = (i == numSamples - 1) ? 1e10 : tValues[i + 1] - tValues[i];
            double alpha = 1.0 - Math.exp(-density * delta);
            
            // 累积颜色
            color = color.blend(sampleColor, alpha * transmittance);
            transmittance *= (1.0 - alpha);
        }
        
        return color;
    }

    // 训练模型
    public void train(List<ImageData> images, int epochs, double learningRate) {
        Random random = new Random();
        
        for (int epoch = 0; epoch < epochs; epoch++) {
            double totalLoss = 0.0;
            
            // 随机选择一张图像
            ImageData imageData = images.get(random.nextInt(images.size()));
            
            // 随机选择一个像素
            int x = random.nextInt(imageData.width);
            int y = random.nextInt(imageData.height);
            
            // 生成射线
            Ray ray = generateRay(imageData.cameraPos, imageData.cameraDir, x, y, imageData.width, imageData.height);
            
            // 渲染射线
            Color renderedColor = renderRay(ray);
            
            // 获取真实颜色
            Color targetColor = imageData.getColor(x, y);
            
            // 计算损失
            double loss = Math.pow(renderedColor.r - targetColor.r, 2) +
                         Math.pow(renderedColor.g - targetColor.g, 2) +
                         Math.pow(renderedColor.b - targetColor.b, 2);
            
            totalLoss += loss;
            
            // 这里应该有反向传播和参数更新的代码
            // 简化实现，实际需要计算梯度并更新网络参数
            
            if (epoch % 100 == 0) {
                System.out.println("Epoch " + epoch + ", Loss: " + loss);
            }
        }
    }

    // 生成射线
    private Ray generateRay(Vector3 cameraPos, Vector3 cameraDir, int x, int y, int width, int height) {
        // 简化的相机模型
        double aspectRatio = (double) width / height;
        double fov = Math.toRadians(90);
        double pixelSize = 2.0 * Math.tan(fov / 2) / height;
        
        // 计算像素在相机坐标系中的位置
        double u = (x - width / 2.0) * pixelSize;
        double v = (height / 2.0 - y) * pixelSize;
        
        // 计算射线方向
        Vector3 up = new Vector3(0, 1, 0);
        Vector3 right = cameraDir.cross(up).normalize();
        up = right.cross(cameraDir).normalize();
        
        Vector3 direction = cameraDir.add(right.multiply(u)).add(up.multiply(v)).normalize();
        
        return new Ray(cameraPos, direction);
    }
}

// 图像数据类
class ImageData {
    Color[][] pixels;
    int width, height;
    Vector3 cameraPos;
    Vector3 cameraDir;

    public ImageData(int width, int height, Vector3 cameraPos, Vector3 cameraDir) {
        this.width = width;
        this.height = height;
        this.cameraPos = cameraPos;
        this.cameraDir = cameraDir;
        this.pixels = new Color[width][height];
    }

    public void setColor(int x, int y, Color color) {
        pixels[x][y] = color;
    }

    public Color getColor(int x, int y) {
        return pixels[x][y];
    }

    // 从文件加载图像
    public static ImageData loadImage(String filePath, Vector3 cameraPos, Vector3 cameraDir) {
        try {
            BufferedImage image = ImageIO.read(new File(filePath));
            int width = image.getWidth();
            int height = image.getHeight();
            
            ImageData data = new ImageData(width, height, cameraPos, cameraDir);
            
            for (int y = 0; y < height; y++) {
                for (int x = 0; x < width; x++) {
                    int rgb = image.getRGB(x, y);
                    int r = (rgb >> 16) & 0xFF;
                    int g = (rgb >> 8) & 0xFF;
                    int b = rgb & 0xFF;
                    
                    data.setColor(x, y, new Color(r / 255.0, g / 255.0, b / 255.0));
                }
            }
            
            return data;
        } catch (IOException e) {
            e.printStackTrace();
            return null;
        }
    }

    // 保存图像
    public void saveImage(String filePath) {
        try {
            BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
            
            for (int y = 0; y < height; y++) {
                for (int x = 0; x < width; x++) {
                    image.setRGB(x, y, pixels[x][y].toRGB());
                }
            }
            
            ImageIO.write(image, "png", new File(filePath));
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}

// 主类
public class NeRFDemo {
    public static void main(String[] args) {
        // 创建一个简单的场景
        List<ImageData> trainingImages = new ArrayList<>();
        
        // 这里应该加载实际的训练图像
        // 简化示例，创建一个虚拟场景
        
        int width = 100;
        int height = 100;
        
        // 从不同角度观察的图像
        for (int angle = 0; angle < 360; angle += 30) {
            double radians = Math.toRadians(angle);
            Vector3 cameraPos = new Vector3(Math.cos(radians) * 5, 0, Math.sin(radians) * 5);
            Vector3 cameraDir = new Vector3(-Math.cos(radians), 0, -Math.sin(radians));
            
            ImageData image = new ImageData(width, height, cameraPos, cameraDir);
            
            // 填充一些简单的图案作为示例
            for (int y = 0; y < height; y++) {
                for (int x = 0; x < width; x++) {
                    double u = (x - width / 2.0) / (width / 2.0);
                    double v = (y - height / 2.0) / (height / 2.0);
                    
                    double distance = Math.sqrt(u * u + v * v);
                    if (distance < 0.5) {
                        image.setColor(x, y, new Color(1, 0, 0)); // 红色圆形
                    } else {
                        image.setColor(x, y, new Color(0, 0, 0)); // 黑色背景
                    }
                }
            }
            
            trainingImages.add(image);
        }
        
        // 创建NeRF模型
        NeRF nerf = new NeRF(64, 2.0, 8.0);
        
        // 训练模型
        System.out.println("开始训练NeRF模型...");
        nerf.train(trainingImages, 1000, 0.001);
        
        // 渲染一个新视角
        Vector3 newCameraPos = new Vector3(5, 0, 0);
        Vector3 newCameraDir = new Vector3(-1, 0, 0);
        
        ImageData renderedImage = new ImageData(width, height, newCameraPos, newCameraDir);
        
        System.out.println("渲染新视角图像...");
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                Ray ray = nerf.generateRay(newCameraPos, newCameraDir, x, y, width, height);
                Color color = nerf.renderRay(ray);
                renderedImage.setColor(x, y, color);
            }
        }
        
        // 保存渲染结果
        renderedImage.saveImage("nerf_render.png");
        System.out.println("渲染完成，图像已保存为 nerf_render.png");
    }
}