Upload 5 files

.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
+assets/example_depth.png filter=lfs diff=lfs merge=lfs -text
+assets/example_rgb.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,9 @@
+# Geometry Prior Visualization Demo
+
+This demo shows the visualization of geometry priors from RGB and depth images.
+
+## Usage
+Upload an RGB image and a depth image to see the visualization result.
+
+## Examples
+The demo includes example images from the NYUDepthv2 dataset.

app.py

@@ -0,0 +1,383 @@
+import torch
+import torch.nn as nn
+
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+from torchvision.transforms.functional import to_pil_image
+from torchvision.transforms import Resize
+import cv2
+import numpy as np
+from torchcam.utils import overlay_mask
+import gradio as gr
+import os
+
+class GeoPrior(nn.Module):
+
+    def __init__(self, embed_dim=128, num_heads=4, initial_value=2, heads_range=6):
+        super().__init__()
+        angle = 1.0 / (10000 ** torch.linspace(0, 1, embed_dim // num_heads // 2))
+        angle = angle.unsqueeze(-1).repeat(1, 2).flatten()
+        self.initial_value = initial_value  
+        self.heads_range = heads_range 
+        self.num_heads = num_heads
+        decay = torch.log(1 - 2 ** (-initial_value - heads_range * torch.arange(num_heads, dtype=torch.float) / num_heads))
+        self.register_buffer('angle', angle)
+        self.register_buffer('decay', decay)
+        
+    def generate_pos_decay(self, H: int, W: int):
+        '''
+        generate 2d decay mask, the result is (HW)*(HW)
+        '''
+        index_h = torch.arange(H).to(self.decay) #保持一個類型
+        index_w = torch.arange(W).to(self.decay) #
+        grid = torch.meshgrid([index_h, index_w])
+        grid = torch.stack(grid, dim=-1).reshape(H*W, 2) #(H*W 2)
+        mask = grid[:, None, :] - grid[None, :, :] #(H*W H*W 2)
+        mask = (mask.abs()).sum(dim=-1)
+        mask = mask #* self.decay[:, None, None]  #(n H*W H*W)
+        return mask
+    
+    def generate_2d_depth_decay(self, H: int, W: int, depth_grid):
+        '''
+        generate 2d decay mask, the result is (HW)*(HW)
+        '''
+        # index_h = torch.arange(H).to(self.decay) #保持一個類型
+        # index_w = torch.arange(W).to(self.decay) #
+        # grid = torch.meshgrid([index_h, index_w])
+        # grid = torch.stack(grid, dim=-1).reshape(H*W, 2) #(H*W 2)
+        # to do: resize depth_grid to H,W
+        # print(depth_grid.shape,H,W,'2d')
+        B,_,H,W = depth_grid.shape
+        grid_d = depth_grid.reshape(B, H*W, 1)
+        print(grid_d.dtype,'aaaaaaaaaaaaaaaaaa')
+        # exit()
+        mask_d = grid_d[:, :, None, :] - grid_d[:, None,:, :] #(H*W H*W)
+        # mask = grid[:, None, :] - grid[None, :, :] #(H*W H*W 2)
+        # print(mask_d.shape, self.decay[None, :, None, None].shape,'11111')
+        mask_d = (mask_d.abs()).sum(dim=-1)
+        # print(torch.max(mask_d),torch.min(mask_d))
+        # exit()
+        mask_d = mask_d.unsqueeze(1) #* self.decay[None, :, None, None].cpu()  #(n H*W H*W)
+        return mask_d
+    
+    
+    
+    def forward(self, slen, depth_map, activate_recurrent=False, chunkwise_recurrent=False):
+        '''
+        slen: (h, w)
+        h * w == l
+        recurrent is not implemented
+        '''
+        # print(depth_map.shape,'depth_map')
+        depth_map = F.interpolate(depth_map, size=slen,mode='bilinear',align_corners=False)
+        # print(depth_map.shape,'downsampled')
+        depth_map = depth_map.float()
+        # depth_map = Resize(slen[0],slen[1])(depth_map).reshape(slen[0],slen[1])
+        
+        index = torch.arange(slen[0]*slen[1]).to(self.decay)
+        sin = torch.sin(index[:, None] * self.angle[None, :]) #(l d1)
+        sin = sin.reshape(slen[0], slen[1], -1) #(h w d1)
+        cos = torch.cos(index[:, None] * self.angle[None, :]) #(l d1)
+        cos = cos.reshape(slen[0], slen[1], -1) #(h w d1)
+        mask_1 = self.generate_pos_decay(slen[0], slen[1]) #(n l l)
+        mask_d = self.generate_2d_depth_decay(slen[0], slen[1], depth_map)
+        print(torch.max(mask_d),torch.min(mask_d),'-2')
+        mask = mask_d#/torch.max(mask_d, dim=0)[0] #mask.cpu() * (2*(1-
+        mask_sum = (0.85*mask_1.cpu()+0.15*mask) * self.decay[:, None, None].cpu()
+        retention_rel_pos = ((sin, cos), mask, mask_1, mask_sum)
+        print(mask.shape,mask_1.shape)
+        # exit()
+
+        return retention_rel_pos
+
+def fangda(mask, in_size=(480//20,640//20), out_size=(480,640)):
+    new_mask = torch.zeros(out_size)
+    ratio_h, ratio_w = out_size[0]//in_size[0], out_size[1]//in_size[1]
+    for i in range(in_size[0]):
+        for j in range(in_size[1]):
+            new_mask[i*ratio_h:(i+1)*ratio_h,j*ratio_w:(j+1)*ratio_w]=mask[i,j]
+    return new_mask
+
+def put_mask(image,mask,color_rgb=None,border_mask=False,color_temp='jet',num_c='',beta=2,fixed_num=None):
+    mask = mask.numpy()
+    image = cv2.resize(image,dsize=(640,480),fx=1,fy=1,interpolation=cv2.INTER_LINEAR)
+    mask = cv2.resize(mask,dsize=(640,480),fx=1,fy=1,interpolation=cv2.INTER_LINEAR)
+    color=np.zeros((1,1,3), dtype=np.uint8)
+    if color_rgb is not None:
+        color[0,0,2],color[0,0,1],color[0,0,0]=color_rgb
+    else:
+        color[0, 0, 2], color[0, 0, 1], color[0, 0, 0]=120,86,87
+    if fixed_num is not None:
+        mask = ((1-mask/255))
+    else:
+        mask=(1-mask/np.max(mask))#*0.5+0.5
+
+
+
+    
+    result = overlay_mask(to_pil_image(image.astype(np.uint8)), to_pil_image(mask), colormap = color_temp, alpha=0.4)
+
+  
+    return np.array(result)
+
+
+def visualize_geometry_prior(RGB_path, Depth_path, index_list=[[584]], cmap_list = ['jet_r'],x=0,y=0):
+    
+    H = 480//20
+    W = 640//20
+    index_num = int(x//20)+int((y//20+1)*32)
+    index_list = [[index_num]]
+    print(index_num)
+    
+    grid_d = cv2.imread(Depth_path,0)
+    # return grid_d
+    grid_d = cv2.resize(grid_d,dsize=(W,H),fx=1,fy=1,interpolation=cv2.INTER_LINEAR)
+
+    grid_d = torch.tensor(grid_d).reshape(1,1,H,W)
+    grid_d_copy=cv2.imread(Depth_path)
+    grid_d_copy = cv2.resize(grid_d_copy,dsize=(640,480),fx=1,fy=1,interpolation=cv2.INTER_LINEAR)
+    grid_d_copy_gray = cv2.imread(Depth_path,0)
+    grid_d_copy_gray = cv2.resize(grid_d_copy_gray,dsize=(640,480),fx=1,fy=1,interpolation=cv2.INTER_LINEAR)
+    print('min max', torch.max(grid_d), torch.min(grid_d))
+    print(grid_d.shape)
+    grid_d=grid_d.cpu()
+
+    respos = GeoPrior()
+    ((sin,cos), depth_map, mask_1, mask_sum) = respos((H,W), grid_d)
+    print(depth_map.shape, mask_1.shape,'-1')
+    print(torch.max(depth_map),torch.min(depth_map))
+
+    # 
+    img_path = RGB_path
+    img = cv2.imread(img_path)
+    img = cv2.resize(img,dsize=(640,480),fx=1,fy=1,interpolation=cv2.INTER_LINEAR)
+
+    grid_d_old = cv2.imread(Depth_path,0)
+    grid_d_old = cv2.resize(grid_d_old,dsize=(W,H),fx=1,fy=1,interpolation=cv2.INTER_LINEAR)
+    grid_d_old = torch.tensor(grid_d_old).reshape(H*W,1)
+    grid_d=grid_d.cpu()
+    mask_d_old = grid_d_old[:, None, :] - grid_d_old[None, :, :] #(H*W H*W 2)
+    mask_d_old = (mask_d_old.abs()).sum(dim=-1)
+    Color_N=255
+    for i in index_list[0]:#range(0,H*W,4):#range(0,H*W,4):#index_list[i_temp]:#range(0,H*W,1): [242,258]:#range
+        for color_temp in cmap_list:
+
+            temp_mask_d = depth_map[0,0,i,:].reshape(H,W).cpu()
+
+            temp_mask = mask_1[i,:].reshape(H,W).cpu()
+            print(torch.max(temp_mask_d),torch.min(temp_mask_d))
+            temp_mask_d_old = mask_d_old[i,:].reshape(H,W).cpu()
+            temp_mask_sum = mask_sum[0,0,i,:].reshape(H,W).cpu()
+            temp_mask_d=torch.nn.functional.normalize(temp_mask_d, p=2.0, dim=1, eps=1e-12, out=None)
+
+            temp_mask_d = 255*(temp_mask_d-torch.min(temp_mask_d))/(torch.max(temp_mask_d)-torch.min(temp_mask_d))
+            
+            temp_mask = 255*((temp_mask-torch.min(temp_mask))/(torch.max(temp_mask)-torch.min(temp_mask)))
+
+            temp_mask_sum = 255*((temp_mask_sum-torch.min(temp_mask_sum))/(torch.max(temp_mask_sum)-torch.min(temp_mask_sum)))
+            gama =0.55
+            temp_mask_d_old = 255*(temp_mask_d_old-torch.min(temp_mask_d_old))/(torch.max(temp_mask_d_old)-torch.min(temp_mask_d_old))
+            a0=put_mask(img,fangda(temp_mask),color_temp=color_temp)
+            jiange = 255*torch.ones(img.shape[0],20)
+            temp_mask_fuse = torch.cat([fangda(temp_mask),jiange,fangda(temp_mask_d),jiange,fangda(gama*temp_mask+(1-gama)*temp_mask_d),jiange,torch.tensor(grid_d_copy_gray)],dim=1)
+            jiange = np.ones((img.shape[0],20, 3)) * 255
+            
+            a2 = put_mask(img, fangda(temp_mask_d),color_temp=color_temp)
+            print(a2.shape)
+            a3 = put_mask(img,fangda(gama*temp_mask+(1-gama)*temp_mask_d),color_temp=color_temp)
+            # image = np.concatenate([img,grid_d_copy,a0,jiange, a2,jiange,a3,jiange] ,axis=1)
+            # print(image.dtype)
+            # print(np.max(image),np.min(image))
+            # cv2.imshow('./temp/demo_'+color_temp+"_"+str(i)+'.png',image.astype(np.uint8))
+    return a3.astype(np.uint8)
+
+# 新增gradio接口函数
+def process_images(rgb_image, depth_image):
+    """
+    处理上传的图像并返回可视化结果
+    
+    Args:
+        rgb_image: gradio上传的RGB图像
+        depth_image: gradio上传的深度图像
+    Returns:
+        可视化结果图像
+    """
+    # 保存临时文件
+    temp_rgb_path = "temp_rgb.jpg"
+    temp_depth_path = "temp_depth.png"
+    
+    # 保存上传的图像
+    cv2.imwrite(temp_rgb_path, cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR))
+    cv2.imwrite(temp_depth_path, depth_image)
+    
+    # 调用原有的可视化函数
+    try:
+        result = visualize_geometry_prior(temp_rgb_path, temp_depth_path,x=x,y=y)
+        
+        # 清理临时文件
+        os.remove(temp_rgb_path)
+        os.remove(temp_depth_path)
+        
+        # 转换颜色空间
+        result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
+        return result
+    except Exception as e:
+        print(f"Error during processing: {str(e)}")
+        return None
+    finally:
+        # 确保临时文件被删除
+        if os.path.exists(temp_rgb_path):
+            os.remove(temp_rgb_path)
+        if os.path.exists(temp_depth_path):
+            os.remove(temp_depth_path)
+
+def draw_star(image, x, y, size=20, color=(255, 0, 0), thickness=2):
+    """在图像上绘制五角星"""
+    # 计算五角星的顶点
+    pts = np.array([[x, y - size],  # 顶部点
+                    [x + size * 0.588, y + size * 0.809],  # 右下
+                    [x - size * 0.951, y - size * 0.309],  # 左上
+                    [x + size * 0.951, y - size * 0.309],  # 右上
+                    [x - size * 0.588, y + size * 0.809]], np.int32)  # 左下
+    
+    # 绘制五角星
+    cv2.polylines(image, [pts], True, color, thickness)
+    return image
+
+# 创建Gradio界面
+def create_demo():
+    with gr.Blocks() as demo:
+        gr.Markdown("# Geometry Prior Visualization Demo")
+        gr.Markdown("""
+        ### Instructions:
+        1. Upload RGB and Depth images
+        2. Enter X (0-640) and Y (0-480) coordinates
+        3. A star marker will be shown on the images at the selected position
+        4. Click "Generate Visualization" to create the visualization
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                rgb_input = gr.Image(label="Upload RGB Image")
+                depth_input = gr.Image(label="Upload Depth Image", image_mode="L")
+                with gr.Row():
+                    x_coord = gr.Number(label="X (0-640)", value=160, minimum=0, maximum=640)
+                    y_coord = gr.Number(label="Y (0-480)", value=270, minimum=0, maximum=480)
+                coordinates_text = gr.Textbox(label="Grid Position and Index", interactive=False)
+            
+            with gr.Column():
+                marked_rgb = gr.Image(label="Marked RGB Image")
+                marked_depth = gr.Image(label="Marked Depth Image")
+                output_image = gr.Image(label="Visualization Result")
+                status_text = gr.Textbox(label="Status", interactive=False)
+
+        def update_coordinates_and_images(rgb_image, depth_image, x, y):
+            # 确保坐标在有效范围内
+            x = max(0, min(640, float(x)))
+            y = max(0, min(480, float(y)))
+            
+            # 计算在24x32网格中的位置
+            H, W = 480//20, 640//20  # 24x32
+            scaled_x = int(x * W / 640)
+            scaled_y = int(y * H / 480)
+            
+            # 计算一维索引
+            grid_index = scaled_y * W + scaled_x
+            
+            # 在RGB图像上绘制五角星
+            rgb_marked = rgb_image.copy()
+            if len(rgb_marked.shape) == 2:  # 如果是灰度图
+                rgb_marked = cv2.cvtColor(rgb_marked, cv2.COLOR_GRAY2BGR)
+            elif rgb_marked.shape[2] == 4:  # 如果是RGBA
+                rgb_marked = cv2.cvtColor(rgb_marked, cv2.COLOR_RGBA2BGR)
+            rgb_marked = draw_star(rgb_marked, int(x), int(y), size=20, color=(255, 0, 0))
+            
+            # 在深度图像上绘制五角星
+            depth_marked = depth_image.copy()
+            if len(depth_marked.shape) == 2:  # 如果是单通道
+                depth_marked = cv2.cvtColor(depth_marked, cv2.COLOR_GRAY2BGR)
+            depth_marked = draw_star(depth_marked, int(x), int(y), size=20, color=(0, 255, 0))
+            
+            return (f"Grid position: ({scaled_x}, {scaled_y}), Index: {grid_index}", 
+                   rgb_marked, 
+                   depth_marked)
+
+        # 坐标更新按钮
+        coord_update_btn = gr.Button("Update Coordinates")
+        coord_update_btn.click(
+            fn=update_coordinates_and_images,
+            inputs=[rgb_input, depth_input, x_coord, y_coord],
+            outputs=[coordinates_text, marked_rgb, marked_depth]
+        )
+        
+        def process_with_status(rgb_image, depth_image, coords_text, x, y):
+            try:
+                # 保存临时文件
+                temp_rgb_path = "temp_rgb.jpg"
+                temp_depth_path = "temp_depth.png"
+                
+                # 保存上传的图像
+                cv2.imwrite(temp_rgb_path, cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR))
+                cv2.imwrite(temp_depth_path, depth_image)
+                
+                if coords_text:
+                    index = int(coords_text.split("Index: ")[-1])
+                    index_list = [[index]]
+                else:
+                    index_list = [[584]]  # 默认值
+                
+                result = visualize_geometry_prior(temp_rgb_path, temp_depth_path, index_list=index_list, x=x, y=y)
+                
+                # 清理临时文件
+                os.remove(temp_rgb_path)
+                os.remove(temp_depth_path)
+                
+                # 转换颜色空间以正确显示
+                result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
+                return result, "Processing completed successfully!"
+            except Exception as e:
+                # 确保临时文件被删除
+                if os.path.exists(temp_rgb_path):
+                    os.remove(temp_rgb_path)
+                if os.path.exists(temp_depth_path):
+                    os.remove(temp_depth_path)
+                return None, f"Error: {str(e)}"
+
+        process_btn = gr.Button("Generate Visualization")
+        process_btn.click(
+            fn=process_with_status,
+            inputs=[rgb_input, depth_input, coordinates_text, x_coord, y_coord],
+            outputs=[output_image, status_text]
+        )
+        
+        # 添加自动更新功能
+        x_coord.change(
+            fn=update_coordinates_and_images,
+            inputs=[rgb_input, depth_input, x_coord, y_coord],
+            outputs=[coordinates_text, marked_rgb, marked_depth]
+        )
+        y_coord.change(
+            fn=update_coordinates_and_images,
+            inputs=[rgb_input, depth_input, x_coord, y_coord],
+            outputs=[coordinates_text, marked_rgb, marked_depth]
+        )
+        
+        gr.Examples(
+            examples=[
+                ["assets/example_rgb.jpg", "assets/example_depth.png"]
+            ],
+            inputs=[rgb_input, depth_input]
+        )
+    
+    return demo
+
+# 启动代码
+if __name__ == "__main__":
+    demo = create_demo()
+    demo.queue()
+    demo.launch(
+        server_name="0.0.0.0",
+        share=True,
+        debug=True
+    )

requirements.txt

@@ -0,0 +1,6 @@
+gradio
+torch
+torchvision
+opencv-python-headless
+numpy
+torchcam