initial commit

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,21 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Models and Engines
+*.onnx
+*.onnx.data
+*.pth
+*.engine
+
+# Images
+*.png
+*.jpeg
+*.JPG
+
+# Videos
+*.mp4
+
+# Logs
+logs/

.huggingface.yaml

@@ -0,0 +1,3 @@
+sdk: gradio
+python_version: '3.12'
+requirements_file: requirements.txt

app.py

@@ -0,0 +1,214 @@
+import gradio as gr
+import torch
+import numpy as np
+import cv2
+import tempfile
+from PIL import Image
+import torchvision.transforms as transforms
+import matplotlib.pyplot as plt
+from models.pvsdnet_model import PVSDNet
+from models.pvsdnet_lite_model import PVSDNet_Lite
+import helperFunctions as helper
+import parameters_pvsdnet as params
+from huggingface_hub import hf_hub_download
+import joblib
+
+REPO_ID = "3ZadeSSG/PVSDNet"
+print("Downloading/Loading checkpoints from Hugging Face Hub...")
+MODEL_PVSDNET_LITE_LOCATION = hf_hub_download(
+    repo_id=REPO_ID, 
+    filename="pvsdnet_lite_model.pth"
+)
+
+MODEL_PVSDNET_LOCATION = hf_hub_download(
+    repo_id=REPO_ID, 
+    filename="pvsdnet_model.pth"
+)
+
+print(f"Large Model loaded at: {MODEL_PVSDNET_LITE_LOCATION}")
+print(f"Lite Model loaded at: {MODEL_PVSDNET_LOCATION}")
+
+DEVICE = params.DEVICE
+
+def getPositionVector(x, y, z):
+    vector = torch.zeros((1, 3), dtype=torch.float)
+    normalized_x = (float(format(x, '.7f')) - (-0.1)) / (0.1 - (-0.1))
+    normalized_y = (float(format(y, '.7f')) - (-0.1)) / (0.1 - (-0.1))
+    normalized_z = (float(format(z, '.7f')) - (-0.1)) / (0.1 - (-0.1))
+    vector[0][0] = normalized_x
+    vector[0][1] = normalized_y
+    vector[0][2] = normalized_z
+    return vector
+
+def generateCircularTrajectory(radius, num_frames):
+    angles = np.linspace(0, 2 * np.pi, num_frames, endpoint=False)
+    return [[radius * np.cos(angle), radius * np.sin(angle), 0] for angle in angles]
+
+def generateSwingTrajectory(radius, num_frames):
+    angles = np.linspace(0, 2 * np.pi, num_frames, endpoint=False)
+    return [[radius * np.cos(angle), 0, radius * np.sin(angle)] for angle in angles]
+
+def create_video_from_memory(frames, fps=30):
+    if not frames:
+        return None
+    height, width, _ = frames[0].shape
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    temp_video = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
+    out = cv2.VideoWriter(temp_video.name, fourcc, fps, (width, height))
+    for frame in frames:
+        out.write(frame)
+    out.release()
+    return temp_video.name
+
+def process_image(img, video_type, radius, num_frames, num_loops, model_type):
+    if img is None:
+        return None, None
+        
+    height, width = 256, 256
+
+    min_dim = min(img.width, img.height)
+    left = (img.width - min_dim) / 2
+    top = (img.height - min_dim) / 2
+    right = (img.width + min_dim) / 2
+    bottom = (img.height + min_dim) / 2
+    img = img.crop((left, top, right, bottom))
+
+    if model_type == "PVSDNet Lite":
+        model = PVSDNet_Lite(total_image_input=params.params_number_input)
+        checkpoint = MODEL_PVSDNET_LITE_LOCATION
+    else:
+        model = PVSDNet(total_image_input=params.params_number_input)
+        checkpoint = MODEL_PVSDNET_LOCATION
+    
+    try:
+        model = helper.load_Checkpoint(checkpoint, model, load_cpu=True)
+    except Exception as e:
+        print(f"Error loading checkpoint {checkpoint}: {e}")
+        
+    model.to(DEVICE)
+    model.eval()
+
+    transform = transforms.Compose([
+        transforms.Resize((height, width)),
+        transforms.ToTensor()
+    ])
+    
+    img_input = img.convert('RGB')
+    img_input = transform(img_input).unsqueeze(0).to(DEVICE)
+
+    if video_type == "Circle":
+        raw_traj = generateCircularTrajectory(radius, num_frames)
+        trajectory = [(p[0], p[1], 0) for p in raw_traj]
+            
+    elif video_type == "Swing":
+        raw_traj = generateSwingTrajectory(radius, num_frames)
+        trajectory = raw_traj
+    else:
+         raw_traj = generateCircularTrajectory(radius, num_frames)
+         trajectory = [(p[0], p[1], 0) for p in raw_traj]
+
+    view_frames = []
+    depth_frames = []
+
+# Run inference for a single loop (trajectory) to save computation
+    for x, y, z in trajectory:
+        pos = getPositionVector(x, y, z).unsqueeze(0).to(DEVICE)
+        
+        with torch.no_grad():
+            predicted_img, predicted_depth = model(img_input, pos)
+        
+        p_img = predicted_img[0].detach().cpu().permute(1, 2, 0).numpy()
+        p_img = np.clip(p_img, 0, 1)
+        p_img = (p_img * 255).astype(np.uint8)
+        p_img_bgr = cv2.cvtColor(p_img, cv2.COLOR_RGB2BGR)
+        view_frames.append(p_img_bgr)
+
+        d_img = predicted_depth.squeeze().detach().cpu().numpy()
+        d_min, d_max = d_img.min(), d_img.max()
+        if d_max - d_min > 1e-6:
+             d_img = (d_img - d_min) / (d_max - d_min)
+        else:
+             d_img = np.zeros_like(d_img)
+        
+        d_img_colored = plt.get_cmap('inferno')(d_img)[:, :, :3]
+        d_img_colored = (d_img_colored * 255).astype(np.uint8)
+        d_img_bgr = cv2.cvtColor(d_img_colored, cv2.COLOR_RGB2BGR)
+        depth_frames.append(d_img_bgr)
+
+    # Repeat the frames for the requested number of loops
+    view_frames = view_frames * int(num_loops)
+    depth_frames = depth_frames * int(num_loops)
+
+    fps = 60
+    view_video_path = create_video_from_memory(view_frames, fps=fps)
+    depth_video_path = create_video_from_memory(depth_frames, fps=fps)
+
+    return view_video_path, depth_video_path
+
+with gr.Blocks(title="PVSDNet: View & Depth Synthesis", theme="default") as demo:
+    gr.Markdown(
+    """
+    ## PVSDNet: Joint Depth Prediction and View Synthesis via Shared Latent Spaces in Real-Time
+    * Upload an image and get a mini video showing capability of novel view and depth synthesis.
+    
+    **Note:** Huggingface demo is running on CPU so inference speeds will be slow. Inference might take around 2 mins.
+    ### Head to our [Project Page](https://realistic3d-miun.github.io/PVSDNet/) for more details about the models.
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            img_input = gr.Image(type="pil", label="Input Image", height=256)
+            
+            with gr.Group():
+                video_type = gr.Dropdown(["Circle", "Swing"], label="Trajectory Type", value="Swing")
+                model_type = gr.Dropdown(["PVSDNet", "PVSDNet Lite"], label="Model Type", value="PVSDNet")
+            
+            with gr.Accordion("Advanced Settings", open=False):
+                radius = gr.Slider(0.01, 0.1, value=0.06, label="Motion Radius")
+                num_frames = gr.Slider(10, 120, value=60, step=1, label="Frames per Loop")
+                num_loops = gr.Slider(1, 6, value=3, step=1, label="Number of Loops")
+            
+            submit_btn = gr.Button("Generate", variant="primary")
+
+        with gr.Column():
+            video_output = gr.Video(label="Generated View Video", height=256)
+            depth_video_output = gr.Video(label="Generated Depth Video", height=256)
+
+    submit_btn.click(
+        fn=process_image,
+        inputs=[img_input, video_type, radius, num_frames, num_loops, model_type],
+        outputs=[video_output, depth_video_output]
+    )
+
+    gr.Markdown("### Example Images: Click to Load")
+    with gr.Column():
+        with gr.Row():
+            sample_1 = gr.Image("./samples/PVSDNet_Samples/COCO_59_source_image.png", label="COCO Example 59", height=150, interactive=False, show_label=True)
+            sample_2 = gr.Image("./samples/PVSDNet_Samples/COCO_16_source_image.png", label="COCO Example 16", height=150, interactive=False, show_label=True)
+            sample_3 = gr.Image("./samples/PVSDNet_Samples/COCO_755_source_image.png", label="COCO Example 755", height=150, interactive=False, show_label=True)
+        
+        with gr.Row():
+            sample_4 = gr.Image("./samples/PVSDNet_Samples/COCO_223_source_image.png", label="COCO Example 223", height=150, interactive=False, show_label=True)
+            sample_5 = gr.Image("./samples/PVSDNet_Samples/COCO_23_source_image.png", label="COCO Example 23", height=150, interactive=False, show_label=True)
+            sample_6 = gr.Image("./samples/PVSDNet_Samples/person.jpeg", label="Person", height=150, interactive=False, show_label=True)
+            
+
+        with gr.Row():
+            sample_7 = gr.Image("./samples/PVSDNet_Samples/flower.png", label="Flower", height=150, interactive=False, show_label=True)
+            sample_8 = gr.Image("./samples/PVSDNet_Samples/person_2.jpeg", label="Person", height=150, interactive=False, show_label=True)
+            sample_9 = gr.Image("./samples/PVSDNet_Samples/bakery.jpeg", label="Bakery", height=150, interactive=False, show_label=True)
+            
+    sample_1.select(fn=lambda: Image.open("./samples/PVSDNet_Samples/COCO_59_source_image.png"), outputs=img_input)
+    sample_2.select(fn=lambda: Image.open("./samples/PVSDNet_Samples/COCO_16_source_image.png"), outputs=img_input)
+    sample_3.select(fn=lambda: Image.open("./samples/PVSDNet_Samples/COCO_755_source_image.png"), outputs=img_input)
+
+    sample_4.select(fn=lambda: Image.open("./samples/PVSDNet_Samples/COCO_223_source_image.png"), outputs=img_input)
+    sample_5.select(fn=lambda: Image.open("./samples/PVSDNet_Samples/COCO_23_source_image.png"), outputs=img_input)
+    sample_6.select(fn=lambda: Image.open("./samples/PVSDNet_Samples/person.jpeg"), outputs=img_input)
+
+    sample_7.select(fn=lambda: Image.open("./samples/PVSDNet_Samples/flower.png"), outputs=img_input)
+    sample_8.select(fn=lambda: Image.open("./samples/PVSDNet_Samples/person_2.jpeg"), outputs=img_input)
+    sample_9.select(fn=lambda: Image.open("./samples/PVSDNet_Samples/bakery.jpeg"), outputs=img_input)
+    
+if __name__ == "__main__":
+    demo.launch()

helperFunctions.py

@@ -0,0 +1,26 @@
+import torch
+import os
+import torch.nn.functional as F
+
+def save_checkpoint(model, filelocation, save_parallel = True):
+    if save_parallel:
+        torch.save(model.module.state_dict(), filelocation)
+    else:
+        torch.save(model.state_dict(), filelocation)
+
+def load_Checkpoint(fileLocation,model, load_cpu=False):
+    if load_cpu:
+        model.load_state_dict(torch.load(fileLocation,map_location=lambda storage, loc: storage))
+    else:
+        model.load_state_dict(torch.load(fileLocation))
+    return model
+
+def writeLog(logList, filename):
+    with open(filename, 'w') as outfile:
+        outfile.write("\n".join(logList))
+
+
+def kl_loss(mu, logvar):
+    return -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).mean()        
+
+

models/pvsdnet_lite_model.py

@@ -0,0 +1,421 @@
+import sys
+import os
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import warnings
+warnings.filterwarnings("ignore")
+import torchvision
+import rff.layers as rff
+import parameters_pvsdnet as params
+import helperFunctions as helper
+
+def getLinearLayer(in_feat, out_feat, activation=nn.ReLU(True)):
+    return nn.Sequential(
+        nn.Linear(in_features=in_feat, out_features=out_feat, bias=True),
+        activation
+    )
+
+def getConvLayer(in_channel,out_channel,stride=1,padding=1,activation=nn.ReLU()):
+    return nn.Sequential(nn.Conv2d(in_channel, 
+                    out_channel,
+                    kernel_size=3,
+                    stride=stride,
+                    padding=padding,
+                    padding_mode='reflect'),
+                    activation)
+
+def getConvTransposeLayer(in_channel, out_channel,kernel=3,stride=1,padding=1,activation=nn.ReLU()):
+    return nn.Sequential(nn.ConvTranspose2d(in_channel,
+                                            out_channel,
+                                            kernel_size = kernel,
+                                            stride=stride,
+                                            padding=padding),
+                                            activation)
+
+class Flatten(nn.Module):
+    def forward(self, input):
+        return input.view(input.size(0), -1)
+
+class UnFlatten(nn.Module):
+    def forward(self, input, size=1):
+        return input.view(input.size(0), 1, params.params_height//8, params.params_width//8)
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.relu = nn.ReLU()
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
+        self.stride = stride
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_channels != out_channels:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels)
+            )
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+
+        out = out + self.shortcut(residual)
+        out = self.relu(out)
+        return out
+
+class MLPEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.m = params.params_m
+        self.positional_encoding = rff.PositionalEncoding(sigma=1,m=self.m)
+        self.layer1 = getLinearLayer(2*3*self.m, 1024) # 2*3*m = 12, here m=32
+        self.dropout1 = nn.Dropout(0.2)
+        self.layer2 = getLinearLayer(1024, 2048)
+        self.dropout2 = nn.Dropout(0.2)
+        self.layer3 = getLinearLayer(2048, (params.params_height//8)*(params.params_width//8))
+        self.unflat = UnFlatten()
+
+        self.up_layer1 = nn.Upsample(scale_factor=2, mode='nearest')
+        self.up_layer2 = nn.Upsample(scale_factor=2, mode='nearest')
+        self.up_layer3 = nn.Upsample(scale_factor=2, mode='nearest')
+    
+
+    def forward(self, x):
+        x = self.positional_encoding(x)
+
+        x = self.layer1(x)
+        x = self.dropout1(x)
+        x = self.layer2(x)
+        x = self.dropout2(x)
+        x = self.layer3(x)
+
+        x = self.unflat(x)
+
+        x = self.up_layer1(x)
+        x = self.up_layer2(x)
+        x = self.up_layer3(x)
+        return x
+    
+class UpperEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        model = torchvision.models.resnet152(pretrained=False)
+        layers = list(model.children())
+        self.ResNetEncoder = torch.nn.Sequential(*layers[:5].copy())
+        del model
+
+    def forward(self, x):
+        x1 = x[:, 0:3, :, :]
+        x1 = self.ResNetEncoder(x1)
+        return x1
+
+    def apply_resnet_encoder(self, x):
+        x1 = x[:, 0:3, :, :]
+        x1 = self.ResNetEncoder(x1)
+        return x1
+
+
+class LowerEncoder(nn.Module):
+    def __init__(self,total_image_input=1):
+        super().__init__()
+        self.encoder_pre = ResidualBlock((total_image_input*3)+1, 20)
+        self.encoder_layer1 = ResidualBlock(20, 30)
+        self.encoder_layer2 = ResidualBlock(30, 50)
+
+        self.encoder_layer3 = nn.Sequential(
+            ResidualBlock(50, 100),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+        
+        self.encoder_layer4 = ResidualBlock(100, 200)
+        self.encoder_layer5 = nn.Sequential(
+            ResidualBlock(200, 200),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+        
+        self.encoder_layer6 = ResidualBlock(200, 200)
+        self.encoder_layer7 = nn.Sequential(
+            ResidualBlock(200, 200),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+        
+        self.encoder_layer8 = ResidualBlock(200, 500)
+        self.encoder_layer9 = nn.Sequential(
+            ResidualBlock(500, 500),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+                                
+        self.encoder_layer10 = ResidualBlock(500, 500)
+        self.encoder_layer11 = ResidualBlock(500, 500)
+        
+    def forward(self, x):
+        x = self.encoder_pre(x)
+        x = self.encoder_layer1(x)
+        x = self.encoder_layer2(x)
+        skip1 = self.encoder_layer3(x)
+        
+        x = self.encoder_layer4(skip1)
+        skip2 = self.encoder_layer5(x)
+
+        x = self.encoder_layer6(skip2)
+        skip3 = self.encoder_layer7(x)
+
+        x = self.encoder_layer8(skip3)
+        skip4 = self.encoder_layer9(x)
+
+        x = self.encoder_layer10(skip4)
+        x = self.encoder_layer11(x)
+
+        return x, [skip1, skip2, skip3, skip4]
+
+class MergeDecoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.decoder_layer1 = ResidualBlock(500, 500)
+        self.decoder_layer2 = ResidualBlock(500, 500)
+        self.decoder_layer3 = ResidualBlock(500, 500)
+        
+        self.decoder_layer4 = nn.Sequential(
+            nn.ConvTranspose2d(500, 200, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer5 = ResidualBlock(200, 200)
+    
+        self.decoder_layer6 = nn.Sequential(
+            nn.ConvTranspose2d(200, 200, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer7 = ResidualBlock(200, 200)
+
+        self.decoder_layer8 = nn.Sequential(
+            nn.ConvTranspose2d(200, 100, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer9 = ResidualBlock(100, 100)
+
+        self.decoder_layer10 = nn.Sequential(
+            nn.ConvTranspose2d(100, 100, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer11 = ResidualBlock(100, 100)
+        self.decoder_layer12 = ResidualBlock(100, 50)
+        self.decoder_layer13 = ResidualBlock(50, 40)
+        self.decoder_layer14 = ResidualBlock(40, 20)
+        self.decoder_layer15 = nn.Sequential(
+            nn.Conv2d(20, 8, 3, stride=1, padding=1),
+            nn.Sigmoid()
+        )
+        self.decoder_layer16 = nn.Sequential(
+            nn.Conv2d(8, 3, 3, stride=1, padding=1),
+            nn.Sigmoid()
+        )
+        
+    def forward(self, x, lower_skip_list, upper_skip_list):
+        x = self.decoder_layer1(x)
+        x = self.decoder_layer2(x)
+        x = x + lower_skip_list[3] + upper_skip_list[1]
+
+        x = self.decoder_layer3(x)
+        x = self.decoder_layer4(x)
+        x = x + lower_skip_list[2] + upper_skip_list[0]
+        
+        x = self.decoder_layer5(x)
+        x = self.decoder_layer6(x)
+        x = x + lower_skip_list[1]
+
+        x = self.decoder_layer7(x)
+        x = self.decoder_layer8(x)
+        x = x + lower_skip_list[0]
+
+        x = self.decoder_layer9(x)
+        x = self.decoder_layer10(x)
+        x = self.decoder_layer11(x)
+        x = self.decoder_layer12(x)
+        x = self.decoder_layer13(x)
+        x = self.decoder_layer14(x)
+        x = self.decoder_layer15(x)
+        x = self.decoder_layer16(x)        
+        return x
+
+
+class DepthDecoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.decoder_layer1 = ResidualBlock(500, 1400)
+        self.decoder_layer2 = ResidualBlock(1400, 1200)
+        self.decoder_layer3 = ResidualBlock(1200, 1000)
+        
+        self.decoder_layer4 = nn.Sequential(
+            nn.ConvTranspose2d(1000, 800, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer5 = ResidualBlock(800, 600)
+    
+        self.decoder_layer6 = nn.Sequential(
+            nn.ConvTranspose2d(600, 400, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer7 = ResidualBlock(400, 200)
+
+        self.decoder_layer8 = nn.Sequential(
+            nn.ConvTranspose2d(200, 100, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer9 = ResidualBlock(100, 100)
+
+        self.decoder_layer10 = nn.Sequential(
+            nn.ConvTranspose2d(100, 100, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer11 = ResidualBlock(100, 100)
+        self.decoder_layer12 = ResidualBlock(100, 50)
+        self.decoder_layer13 = ResidualBlock(50, 40)
+        self.decoder_layer14 = ResidualBlock(40, 20)
+        self.decoder_layer15 = nn.Sequential(
+            nn.Conv2d(20, 8, 3, stride=1, padding=1),
+            nn.ReLU(True)
+        )
+        self.decoder_layer16 = nn.Sequential(
+            nn.Conv2d(8, 1, 3, stride=1, padding=1),
+            nn.ReLU(True)
+        )
+
+        self.up_refinement_0 = ResidualBlock(200, 800)
+        self.up_refinement_1 = ResidualBlock(500, 1200)
+
+        self.low_refinement_1 = ResidualBlock(200, 400)
+        self.low_refinement_2 = ResidualBlock(200, 800)
+        self.low_refinement_3 = ResidualBlock(500, 1200)
+
+        
+        
+    def forward(self, x, lower_skip_list, upper_skip_list):
+        x = self.decoder_layer1(x)
+        x = self.decoder_layer2(x)
+
+        low_skip_3 = self.low_refinement_3(lower_skip_list[3])
+        up_skip_1 = self.up_refinement_1(upper_skip_list[1])
+        x = x + low_skip_3 + up_skip_1
+
+        x = self.decoder_layer3(x)
+        x = self.decoder_layer4(x)
+
+        low_skip_2 = self.low_refinement_2(lower_skip_list[2])
+        up_skip_0 = self.up_refinement_0(upper_skip_list[0])
+        x = x + low_skip_2 + up_skip_0
+        
+        x = self.decoder_layer5(x)
+        x = self.decoder_layer6(x)
+
+        low_skip_1 = self.low_refinement_1(lower_skip_list[1])
+        x = x + low_skip_1
+
+        x = self.decoder_layer7(x)
+        x = self.decoder_layer8(x)
+        x = x + lower_skip_list[0]
+
+        x = self.decoder_layer9(x)
+        x = self.decoder_layer10(x)
+        x = self.decoder_layer11(x)
+        x = self.decoder_layer12(x)
+        x = self.decoder_layer13(x)
+        x = self.decoder_layer14(x)
+        x = self.decoder_layer15(x)
+        x = self.decoder_layer16(x)        
+        return x
+
+
+class PVSNet_Lite(nn.Module):
+    def __init__(self,total_image_input=1):
+        super().__init__()
+        self.target_positional_embedding = MLPEncoder()
+        self.upper_encoder = UpperEncoder()
+        self.lower_encoder = LowerEncoder(total_image_input)
+        self.merge_decoder = MergeDecoder()
+
+        self.upper_encoder_extra_1 = nn.Sequential(
+            ResidualBlock(256, 200),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+        self.upper_encoder_extra_2 = nn.Sequential(
+            ResidualBlock(200, 500),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+        
+    def forward(self, x, pos):
+        target_position_feature = self.target_positional_embedding(pos)
+
+        # First Encoder Branch
+        upper_features_1 = self.upper_encoder.apply_resnet_encoder(x)
+        upper_features_1 = self.upper_encoder_extra_1(upper_features_1)
+        upper_features_2 = self.upper_encoder_extra_2(upper_features_1)
+
+        # Second Encoder Branch
+        stacked_tensor = torch.cat((x,target_position_feature),dim=1)
+        lower_feature, skip_list = self.lower_encoder(stacked_tensor)
+
+        # Decoder
+        merged_feature = self.merge_decoder(lower_feature, skip_list, [upper_features_1, upper_features_2])
+        return merged_feature
+
+
+
+
+class PVSDNet_Lite(nn.Module):
+    def __init__(self,total_image_input=1):
+        super().__init__()
+        self.target_positional_embedding = MLPEncoder()
+        self.upper_encoder = UpperEncoder()
+        self.lower_encoder = LowerEncoder(total_image_input)
+        self.merge_decoder = MergeDecoder()
+        self.depth_decoder = DepthDecoder()
+
+        self.upper_encoder_extra_1 = nn.Sequential(
+            ResidualBlock(256, 200),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+        self.upper_encoder_extra_2 = nn.Sequential(
+            ResidualBlock(200, 200),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+
+        print("Loading pre-trained nvs net")
+        base_net = PVSNet_Lite(total_image_input)
+        #base_net = helper.load_Checkpoint("./checkpoint/checkpoint_init_pvsnet.pth", base_net, load_cpu=True)
+
+        self.target_positional_embedding = base_net.target_positional_embedding
+        self.upper_encoder = base_net.upper_encoder
+        self.lower_encoder = base_net.lower_encoder
+        self.merge_decoder = base_net.merge_decoder
+        self.upper_encoder_extra_1 = base_net.upper_encoder_extra_1
+        self.upper_encoder_extra_2 = base_net.upper_encoder_extra_2
+        del base_net
+        print("Loading pre-trained nvs net: Done")
+
+        
+    def forward(self, x, pos):
+        target_position_feature = self.target_positional_embedding(pos)
+
+        # First Encoder Branch
+        upper_features_1 = self.upper_encoder.apply_resnet_encoder(x)
+        upper_features_1 = self.upper_encoder_extra_1(upper_features_1)
+        upper_features_2 = self.upper_encoder_extra_2(upper_features_1)
+
+        # Second Encoder Branch
+        stacked_tensor = torch.cat((x,target_position_feature),dim=1)
+        lower_feature, skip_list = self.lower_encoder(stacked_tensor)
+
+        # Decoder
+        merged_feature = self.merge_decoder(lower_feature, skip_list, [upper_features_1, upper_features_2])
+
+        # Depth Decoder
+        depth_feature = self.depth_decoder(lower_feature, skip_list, [upper_features_1, upper_features_2])
+        return merged_feature, depth_feature

models/pvsdnet_model.py

@@ -0,0 +1,402 @@
+import sys
+import os
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import warnings
+warnings.filterwarnings("ignore")
+import torchvision
+import rff.layers as rff
+import parameters_pvsdnet as params
+import helperFunctions as helper
+
+def getLinearLayer(in_feat, out_feat, activation=nn.ReLU(True)):
+    return nn.Sequential(
+        nn.Linear(in_features=in_feat, out_features=out_feat, bias=True),
+        activation
+    )
+
+def getConvLayer(in_channel,out_channel,stride=1,padding=1,activation=nn.ReLU()):
+    return nn.Sequential(nn.Conv2d(in_channel, 
+                    out_channel,
+                    kernel_size=3,
+                    stride=stride,
+                    padding=padding,
+                    padding_mode='reflect'),
+                    activation)
+
+def getConvTransposeLayer(in_channel, out_channel,kernel=3,stride=1,padding=1,activation=nn.ReLU()):
+    return nn.Sequential(nn.ConvTranspose2d(in_channel,
+                                            out_channel,
+                                            kernel_size = kernel,
+                                            stride=stride,
+                                            padding=padding),
+                                            activation)
+
+
+class Flatten(nn.Module):
+    def forward(self, input):
+        return input.view(input.size(0), -1)
+
+class UnFlatten(nn.Module):
+    def forward(self, input, size=1):
+        return input.view(input.size(0), 1, params.params_height//8, params.params_width//8)
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.relu = nn.ReLU()
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
+        self.stride = stride
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_channels != out_channels:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels)
+            )
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+
+        out = out + self.shortcut(residual)
+        out = self.relu(out)
+        return out
+
+class MLPEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.m = params.params_m
+        self.positional_encoding = rff.PositionalEncoding(sigma=1,m=self.m)
+        self.layer1 = getLinearLayer(2*3*self.m, 1024) # 2*3*m = 12, here m=32, when training with Rotation data, set it to 2*6*m
+        self.dropout1 = nn.Dropout(0.2)
+        self.layer2 = getLinearLayer(1024, 2048)
+        self.dropout2 = nn.Dropout(0.2)
+        self.layer3 = getLinearLayer(2048, (params.params_height//8)*(params.params_width//8))
+        self.unflat = UnFlatten()
+
+        self.up_layer1 = nn.Upsample(scale_factor=2, mode='nearest')
+        self.up_layer2 = nn.Upsample(scale_factor=2, mode='nearest')
+        self.up_layer3 = nn.Upsample(scale_factor=2, mode='nearest')
+    
+
+    def forward(self, x):
+        x = self.positional_encoding(x)
+
+        x = self.layer1(x)
+        x = self.dropout1(x)
+        x = self.layer2(x)
+        x = self.dropout2(x)
+        x = self.layer3(x)
+
+        x = self.unflat(x)
+
+        x = self.up_layer1(x)
+        x = self.up_layer2(x)
+        x = self.up_layer3(x)
+        return x
+    
+class UpperEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        model = torchvision.models.resnet152(pretrained=False)
+        layers = list(model.children())
+        self.ResNetEncoder = torch.nn.Sequential(*layers[:5].copy())
+        del model
+
+    def forward(self, x):
+        x1 = x[:, 0:3, :, :]
+        x1 = self.ResNetEncoder(x1)
+        return x1
+
+    def apply_resnet_encoder(self, x):
+        x1 = x[:, 0:3, :, :]
+        x1 = self.ResNetEncoder(x1)
+        return x1
+
+class LowerEncoder(nn.Module):
+    def __init__(self,total_image_input=1):
+        super().__init__()
+        self.encoder_pre = ResidualBlock((total_image_input*3)+1, 20)
+        self.encoder_layer1 = ResidualBlock(20, 30)
+        self.encoder_layer2 = ResidualBlock(30, 50)
+
+        self.encoder_layer3 = nn.Sequential(
+            ResidualBlock(50, 100),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+        
+        self.encoder_layer4 = ResidualBlock(100, 200)
+        self.encoder_layer5 = nn.Sequential(
+            ResidualBlock(200, 400),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+        
+        self.encoder_layer6 = ResidualBlock(400, 600)
+        self.encoder_layer7 = nn.Sequential(
+            ResidualBlock(600, 800),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+        
+        self.encoder_layer8 = ResidualBlock(800, 1000)
+        self.encoder_layer9 = nn.Sequential(
+            ResidualBlock(1000, 1200),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+                                
+        self.encoder_layer10 = ResidualBlock(1200, 1400)
+        self.encoder_layer11 = ResidualBlock(1400, 1600)
+        
+    def forward(self, x):
+        x = self.encoder_pre(x)
+        x = self.encoder_layer1(x)
+        x = self.encoder_layer2(x)
+        skip1 = self.encoder_layer3(x)
+        
+        x = self.encoder_layer4(skip1)
+        skip2 = self.encoder_layer5(x)
+
+        x = self.encoder_layer6(skip2)
+        skip3 = self.encoder_layer7(x)
+
+        x = self.encoder_layer8(skip3)
+        skip4 = self.encoder_layer9(x)
+
+        x = self.encoder_layer10(skip4)
+        x = self.encoder_layer11(x)
+
+        return x, [skip1, skip2, skip3, skip4]
+
+class MergeDecoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.decoder_layer1 = ResidualBlock(1600, 1400)
+        self.decoder_layer2 = ResidualBlock(1400, 1200)
+        self.decoder_layer3 = ResidualBlock(1200, 1000)
+        
+        self.decoder_layer4 = nn.Sequential(
+            nn.ConvTranspose2d(1000, 800, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer5 = ResidualBlock(800, 600)
+    
+        self.decoder_layer6 = nn.Sequential(
+            nn.ConvTranspose2d(600, 400, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer7 = ResidualBlock(400, 200)
+
+        self.decoder_layer8 = nn.Sequential(
+            nn.ConvTranspose2d(200, 100, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer9 = ResidualBlock(100, 100)
+
+        self.decoder_layer10 = nn.Sequential(
+            nn.ConvTranspose2d(100, 100, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer11 = ResidualBlock(100, 100)
+        self.decoder_layer12 = ResidualBlock(100, 50)
+        self.decoder_layer13 = ResidualBlock(50, 40)
+        self.decoder_layer14 = ResidualBlock(40, 20)
+        self.decoder_layer15 = nn.Sequential(
+            nn.Conv2d(20, 8, 3, stride=1, padding=1),
+            nn.Sigmoid()
+        )
+        self.decoder_layer16 = nn.Sequential(
+            nn.Conv2d(8, 3, 3, stride=1, padding=1),
+            nn.Sigmoid()
+        )
+        
+    def forward(self, x, lower_skip_list, upper_skip_list):
+        x = self.decoder_layer1(x)
+        x = self.decoder_layer2(x)
+        x = x + lower_skip_list[3] + upper_skip_list[1]
+
+        x = self.decoder_layer3(x)
+        x = self.decoder_layer4(x)
+        x = x + lower_skip_list[2] + upper_skip_list[0]
+        
+        x = self.decoder_layer5(x)
+        x = self.decoder_layer6(x)
+        x = x + lower_skip_list[1]
+
+        x = self.decoder_layer7(x)
+        x = self.decoder_layer8(x)
+        x = x + lower_skip_list[0]
+
+        x = self.decoder_layer9(x)
+        x = self.decoder_layer10(x)
+        x = self.decoder_layer11(x)
+        x = self.decoder_layer12(x)
+        x = self.decoder_layer13(x)
+        x = self.decoder_layer14(x)
+        x = self.decoder_layer15(x)
+        x = self.decoder_layer16(x)        
+        return x
+
+class DepthDecoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.decoder_layer1 = ResidualBlock(1600, 1400)
+        self.decoder_layer2 = ResidualBlock(1400, 1200)
+        self.decoder_layer3 = ResidualBlock(1200, 1000)
+        
+        self.decoder_layer4 = nn.Sequential(
+            nn.ConvTranspose2d(1000, 800, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer5 = ResidualBlock(800, 600)
+    
+        self.decoder_layer6 = nn.Sequential(
+            nn.ConvTranspose2d(600, 400, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer7 = ResidualBlock(400, 200)
+
+        self.decoder_layer8 = nn.Sequential(
+            nn.ConvTranspose2d(200, 100, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer9 = ResidualBlock(100, 100)
+
+        self.decoder_layer10 = nn.Sequential(
+            nn.ConvTranspose2d(100, 100, 2, stride=2, padding=0), 
+            nn.ReLU(True)
+        )
+        self.decoder_layer11 = ResidualBlock(100, 100)
+        self.decoder_layer12 = ResidualBlock(100, 50)
+        self.decoder_layer13 = ResidualBlock(50, 40)
+        self.decoder_layer14 = ResidualBlock(40, 20)
+        self.decoder_layer15 = nn.Sequential(
+            nn.Conv2d(20, 8, 3, stride=1, padding=1),
+            nn.ReLU(True)
+        )
+        self.decoder_layer16 = nn.Sequential(
+            nn.Conv2d(8, 1, 3, stride=1, padding=1),
+            nn.ReLU(True)
+        )
+        
+    def forward(self, x, lower_skip_list, upper_skip_list):
+        x = self.decoder_layer1(x)
+        x = self.decoder_layer2(x)
+        x = x + lower_skip_list[3] + upper_skip_list[1]
+
+        x = self.decoder_layer3(x)
+        x = self.decoder_layer4(x)
+        x = x + lower_skip_list[2] + upper_skip_list[0]
+        
+        x = self.decoder_layer5(x)
+        x = self.decoder_layer6(x)
+        x = x + lower_skip_list[1]
+
+        x = self.decoder_layer7(x)
+        x = self.decoder_layer8(x)
+        x = x + lower_skip_list[0]
+
+        x = self.decoder_layer9(x)
+        x = self.decoder_layer10(x)
+        x = self.decoder_layer11(x)
+        x = self.decoder_layer12(x)
+        x = self.decoder_layer13(x)
+        x = self.decoder_layer14(x)
+        x = self.decoder_layer15(x)
+        x = self.decoder_layer16(x)        
+        return x
+
+
+class PVSNet(nn.Module):
+    def __init__(self,total_image_input=1):
+        super().__init__()
+        self.target_positional_embedding = MLPEncoder()
+        self.upper_encoder = UpperEncoder()
+        self.lower_encoder = LowerEncoder(total_image_input)
+        self.merge_decoder = MergeDecoder()
+
+        self.upper_encoder_extra_1 = nn.Sequential(
+            ResidualBlock(256, 800),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+        self.upper_encoder_extra_2 = nn.Sequential(
+            ResidualBlock(800, 1200),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+        
+    def forward(self, x, pos):
+        target_position_feature = self.target_positional_embedding(pos)
+
+        # First Encoder Branch
+        upper_features_1 = self.upper_encoder.apply_resnet_encoder(x)
+        upper_features_1 = self.upper_encoder_extra_1(upper_features_1)
+        upper_features_2 = self.upper_encoder_extra_2(upper_features_1)
+
+        # Second Encoder Branch
+        stacked_tensor = torch.cat((x,target_position_feature),dim=1)
+        lower_feature, skip_list = self.lower_encoder(stacked_tensor)
+
+        # Decoder
+        merged_feature = self.merge_decoder(lower_feature, skip_list, [upper_features_1, upper_features_2])
+
+        return merged_feature
+
+class PVSDNet(nn.Module):
+    def __init__(self,total_image_input=1):
+        super().__init__()
+        self.target_positional_embedding = MLPEncoder()
+        self.upper_encoder = UpperEncoder()
+        self.lower_encoder = LowerEncoder(total_image_input)
+        self.merge_decoder = MergeDecoder()
+        self.depth_decoder = DepthDecoder()
+
+        self.upper_encoder_extra_1 = nn.Sequential(
+            ResidualBlock(256, 800),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+        self.upper_encoder_extra_2 = nn.Sequential(
+            ResidualBlock(800, 1200),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+
+        print("Loading pre-trained nvs net")
+        base_net = PVSNet(total_image_input)
+        #base_net = helper.load_Checkpoint("./checkpoint/checkpoint_init_pvsnet.pth", base_net, load_cpu=True) #uncomment if you want to use the pre-trained pvsnet to train with lora
+
+        self.target_positional_embedding = base_net.target_positional_embedding
+        self.upper_encoder = base_net.upper_encoder
+        self.lower_encoder = base_net.lower_encoder
+        self.merge_decoder = base_net.merge_decoder
+        self.upper_encoder_extra_1 = base_net.upper_encoder_extra_1
+        self.upper_encoder_extra_2 = base_net.upper_encoder_extra_2
+        del base_net
+        print("Loading pre-trained nvs net: Done")
+
+
+        
+    def forward(self, x, pos):
+        target_position_feature = self.target_positional_embedding(pos)
+
+        # First Encoder Branch
+        upper_features_1 = self.upper_encoder.apply_resnet_encoder(x)
+        upper_features_1 = self.upper_encoder_extra_1(upper_features_1)
+        upper_features_2 = self.upper_encoder_extra_2(upper_features_1)
+
+        # Second Encoder Branch
+        stacked_tensor = torch.cat((x,target_position_feature),dim=1)
+        lower_feature, skip_list = self.lower_encoder(stacked_tensor)
+
+        # Decoder
+        merged_feature = self.merge_decoder(lower_feature, skip_list, [upper_features_1, upper_features_2])
+
+        # Depth Decoder
+        depth_feature = self.depth_decoder(lower_feature, skip_list, [upper_features_1, upper_features_2])
+        return merged_feature, depth_feature

parameters_pvsdnet.py

@@ -0,0 +1,30 @@
+import os
+
+params_height = 256
+params_width = 256
+params_m = 32
+params_number_input = 1
+params_step_size = 5
+params_gamma = 0.2
+
+TRAIN_LOCATION = "./lf_train.txt"
+VALIDATION_LOCATION = "./lf_validate.txt"
+TEST_LOCATION = "./lf_test.txt"
+LOG_FILE_LOCATION = "./logs/training_log_0.txt"
+CHECKPOINT_LOCATION = "./checkpoint/"
+RESUME_CHECKPOINT_LOCATION = "./checkpoint/checkpoint_best.pth"
+START_CHECKPOINT_LOCATION = "./checkpoint/checkpoint_init.pth"
+DEVICE = "cpu"
+
+BATCH_SIZE = 24
+LEARNING_RATE = 0.00001
+NUM_EPOCHS = 50
+START_EPOCH = 0
+T_max = 50
+PRINT_INTERVAL = 20
+
+os.makedirs("./logs",exist_ok=True)
+os.makedirs("./checkpoint",exist_ok=True)
+os.makedirs("./output",exist_ok=True)
+
+

requirements.txt

@@ -0,0 +1,18 @@
+numpy
+tempfile
+torch==2.9.1
+torchvision==0.24.1
+pytorch-msssim==1.0.0
+pytorchvideo==0.1.5
+gradio==6.2.0
+gradio_client==2.0.2
+opencv-python==4.6.0.66
+pillow==10.4.0
+pillow_heif==0.15.0
+matplotlib==3.10.8
+matplotlib-inline==0.1.6
+tqdm==4.65.0
+moviepy==1.0.3
+scikit-image==0.26.0
+scikit-learn==1.8.0
+scipy==1.11.4