970 lines
39 KiB
Plaintext
970 lines
39 KiB
Plaintext
{
|
|
"nbformat": 4,
|
|
"nbformat_minor": 0,
|
|
"metadata": {
|
|
"colab": {
|
|
"private_outputs": true,
|
|
"provenance": [],
|
|
"gpuType": "T4"
|
|
},
|
|
"kernelspec": {
|
|
"name": "python3",
|
|
"display_name": "Python 3"
|
|
},
|
|
"language_info": {
|
|
"name": "python"
|
|
},
|
|
"accelerator": "GPU"
|
|
},
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"!nvidia-smi"
|
|
],
|
|
"metadata": {
|
|
"id": "9t3_pqAO4jjS"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"\n",
|
|
"#@title Load library\n",
|
|
"TORCH_ENABLED = True\n",
|
|
"SCIPY_ENABLED = True\n",
|
|
"\n",
|
|
"from transformers import pipeline\n",
|
|
"import requests\n",
|
|
"import numpy as np\n",
|
|
"from PIL import Image\n",
|
|
"from math import cos, sin\n",
|
|
"import math as m\n",
|
|
"import os\n",
|
|
"import cv2\n",
|
|
"import pandas as pd\n",
|
|
"import gc\n",
|
|
"try:\n",
|
|
" from scipy.spatial import cKDTree\n",
|
|
"except:\n",
|
|
" SCIPY_ENABLED = False\n",
|
|
"try:\n",
|
|
" import torch\n",
|
|
"except:\n",
|
|
" TORCH_ENABLED = False"
|
|
],
|
|
"metadata": {
|
|
"id": "jAwwGXpiQ81b",
|
|
"cellView": "form"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"\n",
|
|
"#@title Define functions\n",
|
|
"\n",
|
|
"def clean():\n",
|
|
" for I in range(0,20):\n",
|
|
" gc.collect()\n",
|
|
" torch.cuda.empty_cache()\n",
|
|
"\n",
|
|
"def export_PLY(points, name_file=\"model.ply\", multiple_files=False, format_ascii=\"ascii\"):\n",
|
|
"\n",
|
|
" # Write PLY header\n",
|
|
" header = [\n",
|
|
" \"ply\",\n",
|
|
" f\"format {(format_ascii)} 1.0\",\n",
|
|
" f\"element vertex {(points.shape[0])}\",\n",
|
|
" \"property float x\",\n",
|
|
" \"property float y\",\n",
|
|
" \"property float z\",\n",
|
|
" \"property uchar red\",\n",
|
|
" \"property uchar green\",\n",
|
|
" \"property uchar blue\",\n",
|
|
" \"end_header\\n\"\n",
|
|
" ]\n",
|
|
"\n",
|
|
" if not(name_file.endswith(\".ply\")):\n",
|
|
" name_file+=\".ply\"\n",
|
|
"\n",
|
|
" if \"ascii\" in format_ascii:\n",
|
|
" header_cons = \"\\n\".join(header)\n",
|
|
" if not(multiple_files):\n",
|
|
" with open(name_file, 'wb') as f:\n",
|
|
" np.savetxt(f, points, fmt='%g %g %g %d %d %d', header=header_cons, comments='')\n",
|
|
" with open(name_file, 'rb+') as fout:\n",
|
|
" fout.seek(-1, os.SEEK_END)\n",
|
|
" fout.truncate()\n",
|
|
" else:\n",
|
|
" def split_array(arr, n):\n",
|
|
" return [arr[i:i+n] for i in range(0, len(arr), n)]\n",
|
|
" for idx,points_ck in enumerate(split_array(points, 2097152)):\n",
|
|
" header = [\n",
|
|
" \"ply\",\n",
|
|
" f\"format {(format_ascii)} 1.0\",\n",
|
|
" f\"element vertex {(points_ck.shape[0])}\",\n",
|
|
" \"property float x\",\n",
|
|
" \"property float y\",\n",
|
|
" \"property float z\",\n",
|
|
" \"property uchar red\",\n",
|
|
" \"property uchar green\",\n",
|
|
" \"property uchar blue\",\n",
|
|
" \"end_header\\n\"\n",
|
|
" ]\n",
|
|
" name_file_t = name_file.split(\".ply\")[0]+\"_\"+str(idx)+\".ply\"\n",
|
|
" with open(name_file_t, 'wb') as f:\n",
|
|
" np.savetxt(f, points_ck, fmt='%g %g %g %d %d %d', header=header_cons, comments='')\n",
|
|
" with open(name_file_t, 'rb+') as fout:\n",
|
|
" fout.seek(-1, os.SEEK_END)\n",
|
|
" fout.truncate()\n",
|
|
"\n",
|
|
" else:\n",
|
|
" xyz = points[:, :3].astype(np.float32)\n",
|
|
" rgb = points[:, 3:6].astype(np.uint8)\n",
|
|
"\n",
|
|
"\n",
|
|
"\n",
|
|
" if (multiple_files):\n",
|
|
" def split_array(arr, n):\n",
|
|
" return [arr[i:i+n] for i in range(0, len(arr), n)]\n",
|
|
" xyz = split_array(xyz, 2097152)\n",
|
|
" rgb = split_array(rgb, 2097152)\n",
|
|
" else:\n",
|
|
" xyz =[xyz]\n",
|
|
" rgb =[rgb]\n",
|
|
"\n",
|
|
" for idx,xyz_ in enumerate(xyz):\n",
|
|
" header = [\n",
|
|
" \"ply\",\n",
|
|
" f\"format {(format_ascii)} 1.0\",\n",
|
|
" f\"element vertex {(xyz_.shape[0])}\",\n",
|
|
" \"property float x\",\n",
|
|
" \"property float y\",\n",
|
|
" \"property float z\",\n",
|
|
" \"property uchar red\",\n",
|
|
" \"property uchar green\",\n",
|
|
" \"property uchar blue\",\n",
|
|
" \"end_header\"\n",
|
|
" ]\n",
|
|
" structured_array = np.zeros(xyz_.shape[0], dtype=[\n",
|
|
" ('x', '<f4'), # Little-endian float32\n",
|
|
" ('y', '<f4'),\n",
|
|
" ('z', '<f4'),\n",
|
|
" ('red', 'u1'), # Unsigned byte (0-255)\n",
|
|
" ('green', 'u1'),\n",
|
|
" ('blue', 'u1')\n",
|
|
" ])\n",
|
|
" rgb_ = rgb[idx]\n",
|
|
" name_file_t = (name_file.split(\".ply\")[0]+\"_\"+str(idx)+\".ply\" if multiple_files else name_file)\n",
|
|
" structured_array['x'] = xyz_[:, 0]\n",
|
|
" structured_array['y'] = xyz_[:, 1]\n",
|
|
" structured_array['z'] = xyz_[:, 2]\n",
|
|
" structured_array['red'] = rgb_[:, 0]\n",
|
|
" structured_array['green'] = rgb_[:, 1]\n",
|
|
" structured_array['blue'] = rgb_[:, 2]\n",
|
|
" with open(name_file_t, 'wb') as f:\n",
|
|
" f.write('\\n'.join(header).encode('utf-8'))\n",
|
|
" f.write(b'\\n') # Header ends with a newline\n",
|
|
" structured_array.tofile(f)\n",
|
|
"\n",
|
|
"def from2Dto3D_vectorized(arrayDepth, arrayPixel, maxZ, quality=1):\n",
|
|
" # Check if the input shapes are compatible\n",
|
|
" assert arrayPixel.shape[:2] == arrayDepth.shape, \"Pixel and Depth map shapes do not match\"\n",
|
|
" W, H = arrayDepth.shape # Original dimensions\n",
|
|
"\n",
|
|
" if quality > 1:\n",
|
|
" # Calculate new dimensions based on quality\n",
|
|
" new_W = (W - 1) * quality + 1\n",
|
|
" new_H = (H - 1) * quality + 1\n",
|
|
"\n",
|
|
" # Generate new coordinates using linspace\n",
|
|
" new_X = np.linspace(0, W-1, new_W)\n",
|
|
" new_Y = np.linspace(0, H-1, new_H)\n",
|
|
"\n",
|
|
" # Create meshgrid for new coordinates\n",
|
|
" xi, yi = np.meshgrid(new_X, new_Y, indexing='ij') # shapes (new_W, new_H)\n",
|
|
"\n",
|
|
" # Prepare indices for interpolation\n",
|
|
" x0 = np.floor(xi).astype(int)\n",
|
|
" x1 = np.clip(x0 + 1, 0, W-1)\n",
|
|
" y0 = np.floor(yi).astype(int)\n",
|
|
" y1 = np.clip(y0 + 1, 0, H-1)\n",
|
|
"\n",
|
|
" dx = xi - x0\n",
|
|
" dy = yi - y0\n",
|
|
"\n",
|
|
" # Interpolate arrayDepth\n",
|
|
" tl = arrayDepth[x0, y0]\n",
|
|
" tr = arrayDepth[x0, y1]\n",
|
|
" bl = arrayDepth[x1, y0]\n",
|
|
" br = arrayDepth[x1, y1]\n",
|
|
" interpolated_depth = (\n",
|
|
" (1 - dx) * (1 - dy) * tl +\n",
|
|
" dx * (1 - dy) * tr +\n",
|
|
" (1 - dx) * dy * bl +\n",
|
|
" dx * dy * br\n",
|
|
" )\n",
|
|
"\n",
|
|
" # Interpolate arrayPixel for each channel\n",
|
|
" interpolated_pixel = np.zeros((new_W, new_H, 3), dtype=arrayPixel.dtype)\n",
|
|
" for c in range(3):\n",
|
|
" tl_c = arrayPixel[x0, y0, c]\n",
|
|
" tr_c = arrayPixel[x0, y1, c]\n",
|
|
" bl_c = arrayPixel[x1, y0, c]\n",
|
|
" br_c = arrayPixel[x1, y1, c]\n",
|
|
" interpolated_c = (\n",
|
|
" (1 - dx) * (1 - dy) * tl_c +\n",
|
|
" dx * (1 - dy) * tr_c +\n",
|
|
" (1 - dx) * dy * bl_c +\n",
|
|
" dx * dy * br_c\n",
|
|
" )\n",
|
|
" interpolated_pixel[:, :, c] = interpolated_c.astype(arrayPixel.dtype)\n",
|
|
"\n",
|
|
" # Update variables to use interpolated arrays\n",
|
|
" W, H = new_W, new_H\n",
|
|
" arrayDepth = interpolated_depth\n",
|
|
" arrayPixel = interpolated_pixel\n",
|
|
" # Use new_X and new_Y for coordinates\n",
|
|
" final_X, final_Y = np.meshgrid(new_X, new_Y, indexing='ij')\n",
|
|
" else:\n",
|
|
" # Original coordinates\n",
|
|
" x = np.arange(W)\n",
|
|
" y = np.arange(H)\n",
|
|
" final_X, final_Y = np.meshgrid(x, y, indexing='ij') # shapes (W, H)\n",
|
|
"\n",
|
|
" # Calculate scaled Z values\n",
|
|
" d_min = np.min(arrayDepth)\n",
|
|
" d_max = np.max(arrayDepth)\n",
|
|
" if d_max - d_min > 1e-9:\n",
|
|
" Z = -(arrayDepth - d_min) / (d_max - d_min) * maxZ\n",
|
|
" else:\n",
|
|
" Z = np.zeros_like(arrayDepth, dtype=np.float64)\n",
|
|
"\n",
|
|
" # Extract RGB components\n",
|
|
" R = arrayPixel[:, :, 0]\n",
|
|
" G = arrayPixel[:, :, 1]\n",
|
|
" B = arrayPixel[:, :, 2]\n",
|
|
"\n",
|
|
" # Stack all components and reshape to Nx6\n",
|
|
" # Transpose to shape (H, W) for correct ordering when using 'ij' indexing\n",
|
|
" points = np.stack([\n",
|
|
" final_Y.T, final_X.T, Z.T,\n",
|
|
" R.T, G.T, B.T\n",
|
|
" ], axis=-1)\n",
|
|
" points = points.reshape(-1, 6)\n",
|
|
"\n",
|
|
" return points\n",
|
|
"\n",
|
|
"\n",
|
|
"def from2Dto3D_vectorized_torch(arrayDepth, arrayPixel, maxZ, quality=1):\n",
|
|
" # Ensure inputs are PyTorch tensors\n",
|
|
" assert isinstance(arrayDepth, torch.Tensor), \"arrayDepth must be a torch.Tensor\"\n",
|
|
" assert isinstance(arrayPixel, torch.Tensor), \"arrayPixel must be a torch.Tensor\"\n",
|
|
"\n",
|
|
" # Check shape compatibility\n",
|
|
"\n",
|
|
" assert arrayPixel.shape[:2] == arrayDepth.shape, \"Pixel and Depth map shapes do not match\"\n",
|
|
" W, H = arrayDepth.shape # Original dimensions\n",
|
|
"\n",
|
|
" if quality > 1:\n",
|
|
" # Calculate new dimensions\n",
|
|
" new_W = (W - 1) * quality + 1\n",
|
|
" new_H = (H - 1) * quality + 1\n",
|
|
"\n",
|
|
" # Generate coordinates (CHANGED: torch.linspace instead of np.linspace)\n",
|
|
" new_X = torch.linspace(0, W-1, new_W, device=arrayDepth.device)\n",
|
|
" new_Y = torch.linspace(0, H-1, new_H, device=arrayDepth.device)\n",
|
|
"\n",
|
|
" # Create meshgrid (CHANGED: torch.meshgrid with indexing='ij')\n",
|
|
" xi, yi = torch.meshgrid(new_X, new_Y, indexing='ij')\n",
|
|
"\n",
|
|
" # Prepare indices (CHANGED: torch.floor/int/clamp)\n",
|
|
" x0 = torch.floor(xi).long()\n",
|
|
" x1 = torch.clamp(x0 + 1, 0, W-1)\n",
|
|
" y0 = torch.floor(yi).long()\n",
|
|
" y1 = torch.clamp(y0 + 1, 0, H-1)\n",
|
|
"\n",
|
|
" dx = xi - x0\n",
|
|
" dy = yi - y0\n",
|
|
"\n",
|
|
" # Interpolate depth (CHANGED: Tensor indexing)\n",
|
|
" tl = arrayDepth[x0, y0]\n",
|
|
" tr = arrayDepth[x0, y1]\n",
|
|
" bl = arrayDepth[x1, y0]\n",
|
|
" br = arrayDepth[x1, y1]\n",
|
|
" interpolated_depth = (\n",
|
|
" (1 - dx) * (1 - dy) * tl +\n",
|
|
" dx * (1 - dy) * tr +\n",
|
|
" (1 - dx) * dy * bl +\n",
|
|
" dx * dy * br\n",
|
|
" )\n",
|
|
"\n",
|
|
" # Interpolate pixel (CHANGED: Tensor operations preserve dtype)\n",
|
|
" interpolated_pixel = torch.zeros((new_W, new_H, 3),\n",
|
|
" dtype=arrayPixel.dtype,\n",
|
|
" device=arrayPixel.device)\n",
|
|
" for c in range(3):\n",
|
|
" tl_c = arrayPixel[x0, y0, c]\n",
|
|
" tr_c = arrayPixel[x0, y1, c]\n",
|
|
" bl_c = arrayPixel[x1, y0, c]\n",
|
|
" br_c = arrayPixel[x1, y1, c]\n",
|
|
" interpolated_c = (\n",
|
|
" (1 - dx) * (1 - dy) * tl_c +\n",
|
|
" dx * (1 - dy) * tr_c +\n",
|
|
" (1 - dx) * dy * bl_c +\n",
|
|
" dx * dy * br_c\n",
|
|
" )\n",
|
|
" interpolated_pixel[:, :, c] = interpolated_c.to(arrayPixel.dtype)\n",
|
|
"\n",
|
|
" W, H = new_W, new_H\n",
|
|
" arrayDepth = interpolated_depth\n",
|
|
" arrayPixel = interpolated_pixel\n",
|
|
" final_X, final_Y = xi, yi\n",
|
|
" else:\n",
|
|
" # Original coordinates (CHANGED: torch.arange)\n",
|
|
" x = torch.arange(W, device=arrayDepth.device)\n",
|
|
" y = torch.arange(H, device=arrayDepth.device)\n",
|
|
" final_X, final_Y = torch.meshgrid(x, y, indexing='ij')\n",
|
|
"\n",
|
|
" # Calculate Z (CHANGED: Tensor operations)\n",
|
|
" d_min = torch.min(arrayDepth)\n",
|
|
" d_max = torch.max(arrayDepth)\n",
|
|
" if (d_max - d_min).item() > 1e-9: # CHANGED: .item() for scalar comparison\n",
|
|
" Z = -(arrayDepth - d_min) / (d_max - d_min) * maxZ\n",
|
|
" else:\n",
|
|
" Z = torch.zeros_like(arrayDepth)\n",
|
|
"\n",
|
|
" # Extract RGB (CHANGED: Tensor slicing)\n",
|
|
" R = arrayPixel[:, :, 0]\n",
|
|
" G = arrayPixel[:, :, 1]\n",
|
|
" B = arrayPixel[:, :, 2]\n",
|
|
"\n",
|
|
" # Stack components (CHANGED: permute instead of transpose)\n",
|
|
" points = torch.stack([\n",
|
|
" final_Y.permute(1, 0), # Equivalent to .T in NumPy\n",
|
|
" final_X.permute(1, 0),\n",
|
|
" Z.permute(1, 0),\n",
|
|
" R.permute(1, 0),\n",
|
|
" G.permute(1, 0),\n",
|
|
" B.permute(1, 0)\n",
|
|
" ], dim=-1)\n",
|
|
"\n",
|
|
" return points.view(-1, 6) # CHANGED: view instead of reshape\n",
|
|
"\n",
|
|
"def transform_points(points, translate=(0, 0, 0), rotate=(0, 0, 0), scale=(1,1,1)):\n",
|
|
" \"\"\"Apply 3D transformations to points (rotation first, then translation)\"\"\"\n",
|
|
" # Convert rotation angles to radians if needed (modify if using degrees)\n",
|
|
" rx, ry, rz = rotate\n",
|
|
"\n",
|
|
"\n",
|
|
" # Create rotation matrices\n",
|
|
" # X-axis rotation\n",
|
|
" rot_x = np.array([\n",
|
|
" [1, 0, 0],\n",
|
|
" [0, cos(rx), -sin(rx)],\n",
|
|
" [0, sin(rx), cos(rx)]\n",
|
|
" ], dtype=np.float64)\n",
|
|
"\n",
|
|
" # Y-axis rotation\n",
|
|
" rot_y = np.array([\n",
|
|
" [cos(ry), 0, sin(ry)],\n",
|
|
" [0, 1, 0],\n",
|
|
" [-sin(ry), 0, cos(ry)]\n",
|
|
" ], dtype=np.float64)\n",
|
|
"\n",
|
|
" # Z-axis rotation\n",
|
|
" rot_z = np.array([\n",
|
|
" [cos(rz), -sin(rz), 0],\n",
|
|
" [sin(rz), cos(rz), 0],\n",
|
|
" [0, 0, 1]\n",
|
|
" ], dtype=np.float64)\n",
|
|
"\n",
|
|
"\n",
|
|
" # Combined rotation matrix (Z-Y-X order)\n",
|
|
" rotation_matrix = rot_z @ rot_y @ rot_x\n",
|
|
"\n",
|
|
" points[:,0:3] = (points[:,0:3] * np.array(scale, dtype=np.float64)).astype(np.float64)\n",
|
|
"\n",
|
|
" # Apply rotation\n",
|
|
" points[:,0:3] = (points[:,0:3] @ rotation_matrix.T).astype(np.float64)\n",
|
|
"\n",
|
|
" # Apply translation\n",
|
|
" points[:,0:3] = (points[:,0:3] + np.array(translate, dtype=np.float64)).astype(np.float64)\n",
|
|
"\n",
|
|
" return points\n",
|
|
"\n",
|
|
"def project_points(points, fov=60, rotation=None, translation=None, scale=None, aspect_ratio=1, img_size=(1000,1000)):\n",
|
|
" \"\"\"Fast 3D->2D projection with perspective correction\"\"\"\n",
|
|
" if rotation is None:\n",
|
|
" rotation = np.zeros(3)\n",
|
|
" if translation is None:\n",
|
|
" translation = np.zeros(3)\n",
|
|
" if scale is None:\n",
|
|
" scale = np.ones(3)\n",
|
|
"\n",
|
|
" # Calcolo parametri camera\n",
|
|
" fov_rad = m.radians(fov)\n",
|
|
" focal_length = img_size[1] / (2 * m.tan(fov_rad / 2)) # Focal length verticale\n",
|
|
"\n",
|
|
" # Trasformazioni (camera guarda lungo -Z)\n",
|
|
"\n",
|
|
" points = transform_points(points, translate=translation, rotate=rotation, scale=scale)\n",
|
|
"\n",
|
|
" # Seleziona punti DAVANTI alla camera (z < 0)\n",
|
|
" valid = points[:, 2] < 0\n",
|
|
" points_t = points #[valid]\n",
|
|
"\n",
|
|
" if len(points_t) == 0:\n",
|
|
" return np.empty((0, 2)), np.array([]), np.empty((0, 3))\n",
|
|
"\n",
|
|
" # Calcolo coordinate proiettate\n",
|
|
" z = -points_t[:, 2] # Converti in distanza positiva\n",
|
|
" x_proj = (points_t[:, 0] * focal_length) / (z * aspect_ratio)\n",
|
|
" y_proj = (points_t[:, 1] * focal_length) / z\n",
|
|
"\n",
|
|
" # Conversione a coordinate immagine (Y non invertito)\n",
|
|
" pixel_x = (x_proj + img_size[0]/2).astype(int)\n",
|
|
" pixel_y = (y_proj + img_size[1]/2).astype(int) # Rimosso il segno negativo\n",
|
|
"\n",
|
|
" # Clip e normalizzazione\n",
|
|
" pixel_x = np.clip(pixel_x, 0, img_size[0]-1)\n",
|
|
" pixel_y = np.clip(pixel_y, 0, img_size[1]-1)\n",
|
|
" z_norm = (z - z.min()) / (z.max() - z.min())\n",
|
|
"\n",
|
|
" return np.column_stack([pixel_x, pixel_y]), z_norm, points_t[:, 3:6]\n",
|
|
"\n",
|
|
"def project_points_ortho(points, rotation=None, translation=None, scale=None):\n",
|
|
" \"\"\"Orthographic 3D->2D projection\"\"\"\n",
|
|
" if rotation is None:\n",
|
|
" rotation = np.zeros(3)\n",
|
|
" if translation is None:\n",
|
|
" translation = np.zeros(3)\n",
|
|
" if scale is None:\n",
|
|
" scale = np.ones(3)\n",
|
|
"\n",
|
|
"\n",
|
|
" # Apply transformations\n",
|
|
" points = transform_points(points, translate=translation, rotate=rotation, scale=scale)\n",
|
|
"\n",
|
|
" # Orthographic projection\n",
|
|
" x_proj = points[:, 0] # Simple scaling + centering offset\n",
|
|
" y_proj = points[:, 1]\n",
|
|
"\n",
|
|
" # Normalize depth (similar to perspective version)\n",
|
|
" z = points[:, 2]\n",
|
|
" z_min = z.min()\n",
|
|
" z_max = z.max()\n",
|
|
" z_range = z_max - z_min + 1e-5 # Avoid division by zero\n",
|
|
" z_norm = (z - z_min) / z_range\n",
|
|
"\n",
|
|
" colors = points[:,3:6]\n",
|
|
"\n",
|
|
" return np.column_stack([x_proj, y_proj]), z_norm, colors\n",
|
|
"\n",
|
|
"def render_points_fast(points, rotation=None, translation=None, scale=None,\n",
|
|
" img_size=(1000, 1000), color=False, cameraType=0, fov=60, correction=False, ksize=(2,2), thresh=2):\n",
|
|
" \"\"\"Ultra-fast rendering using pure NumPy and PIL\"\"\"\n",
|
|
" # Project points to 2D\n",
|
|
" if cameraType==0:\n",
|
|
" proj, depth, colors = project_points_ortho(points, rotation=rotation, translation=translation)\n",
|
|
" else:\n",
|
|
" proj, depth, colors = project_points(points, fov=fov, aspect_ratio=img_size[1]/img_size[1], rotation=rotation, translation=translation, scale=scale, img_size=img_size)\n",
|
|
"\n",
|
|
" if isinstance(proj, torch.Tensor):\n",
|
|
" proj=proj.detach().cpu().numpy()\n",
|
|
" depth=depth.detach().cpu().numpy()\n",
|
|
" colors=colors.detach().cpu().numpy()\n",
|
|
"\n",
|
|
" # Convert to integer coordinates and filter valid points\n",
|
|
" proj = proj.astype(np.int32)\n",
|
|
" valid = (proj[:, 0] >= 0) & (proj[:, 0] < img_size[0]) & \\\n",
|
|
" (proj[:, 1] >= 0) & (proj[:, 1] < img_size[1])\n",
|
|
"\n",
|
|
" proj = proj[valid]\n",
|
|
" depth = depth[valid]\n",
|
|
" colors = colors[valid]\n",
|
|
"\n",
|
|
" # Grayscale rendering based on depth\n",
|
|
" max_depth = depth.max() if depth.size > 0 else 1\n",
|
|
" min_depth = depth.min() if depth.size > 0 else 0\n",
|
|
" if max_depth - min_depth > 0:\n",
|
|
" depth_normalized = (depth - min_depth) / (max_depth - min_depth)\n",
|
|
" else:\n",
|
|
" depth_normalized = np.zeros_like(depth)\n",
|
|
" intensities = (255 - (depth_normalized * 255)).astype(np.uint16)\n",
|
|
"\n",
|
|
" # Sort points by descending intensity to prioritize closer points\n",
|
|
" sorted_indices = np.argsort(-intensities)\n",
|
|
" proj_sorted = proj[sorted_indices]\n",
|
|
" y_coords = proj_sorted[:, 1]\n",
|
|
" x_coords = proj_sorted[:, 0]\n",
|
|
"\n",
|
|
" df = pd.DataFrame({'y': y_coords, 'x': x_coords})\n",
|
|
" unique_indices = df.drop_duplicates().index.to_numpy()\n",
|
|
"\n",
|
|
" # Extract unique coordinates and colors\n",
|
|
" unique_y = y_coords[unique_indices]\n",
|
|
" unique_x = x_coords[unique_indices]\n",
|
|
"\n",
|
|
" if color:\n",
|
|
" # Extract colors from valid points (assuming points are Nx[x,y,z,r,g,b])\n",
|
|
" point_colors = colors.astype(np.uint16)\n",
|
|
" colors_sorted = point_colors[sorted_indices]\n",
|
|
"\n",
|
|
" unique_colors = colors_sorted[unique_indices]\n",
|
|
"\n",
|
|
" # Create color buffer and assign colors\n",
|
|
" color_buffer = np.ones((img_size[1], img_size[0], 3))\n",
|
|
" color_buffer = color_buffer*-1\n",
|
|
" color_buffer[unique_y, unique_x] = unique_colors\n",
|
|
"\n",
|
|
" img = color_buffer\n",
|
|
" else:\n",
|
|
" point_colors = intensities\n",
|
|
" colors_sorted = point_colors[sorted_indices]\n",
|
|
"\n",
|
|
" unique_colors = colors_sorted[unique_indices]\n",
|
|
"\n",
|
|
" # Create color buffer and assign colors\n",
|
|
" color_buffer = np.ones((img_size[1], img_size[0]))\n",
|
|
" color_buffer = color_buffer*-1\n",
|
|
" color_buffer[unique_y, unique_x] = unique_colors\n",
|
|
"\n",
|
|
" img = color_buffer\n",
|
|
"\n",
|
|
"\n",
|
|
" if correction:\n",
|
|
" point_colors = intensities\n",
|
|
" colors_sorted = point_colors[sorted_indices]\n",
|
|
"\n",
|
|
" unique_colors = colors_sorted[unique_indices]\n",
|
|
"\n",
|
|
" # Create color buffer and assign colors\n",
|
|
" depth_buffer = np.zeros((img_size[1], img_size[0]))\n",
|
|
" depth_buffer[unique_y, unique_x] = unique_colors\n",
|
|
"\n",
|
|
" img_depth = depth_buffer\n",
|
|
"\n",
|
|
" img = np.array(img)\n",
|
|
" img_dep = np.array(img_depth)\n",
|
|
"\n",
|
|
" img_blur = blur_image_excluding_black(img, ksize, thresh)\n",
|
|
" #print(np.unique(img_blur))\n",
|
|
" mask = img_dep < thresh\n",
|
|
"\n",
|
|
" (img[mask]) = (img_blur[mask])\n",
|
|
"\n",
|
|
" img = img\n",
|
|
"\n",
|
|
" return img\n",
|
|
"\n",
|
|
"def clean_points(points, k, m):\n",
|
|
" points = np.asarray(points)\n",
|
|
" if points.size == 0:\n",
|
|
" return points.copy()\n",
|
|
" if k <= 0:\n",
|
|
" raise ValueError(\"k must be a positive integer\")\n",
|
|
" n = len(points)\n",
|
|
" if k > n - 1:\n",
|
|
" return np.empty((0, points.shape[1]))\n",
|
|
" tree = cKDTree(points, balanced_tree=False)\n",
|
|
" # Query k+1 to include the point itself, then select the k-th neighbor\n",
|
|
" distances, _ = tree.query(points, k=k+1, workers=-1)\n",
|
|
" kth_distances = distances[:, k] # k-th neighbor after excluding self\n",
|
|
" mask = kth_distances <= m\n",
|
|
" return points[mask]\n",
|
|
"\n",
|
|
"def interpolate_points(points, alpha=0.5, n_c=3):\n",
|
|
" \"\"\"\n",
|
|
" Interpolate three new points for each original point, positioned between the original and its three nearest neighbors.\n",
|
|
"\n",
|
|
" Parameters:\n",
|
|
" points (numpy.ndarray): Input array of shape (N, 6) where each row is (x, y, z, r, g, b).\n",
|
|
" alpha (float): Interpolation factor (0.0 = original point, 1.0 = neighbor). Default is 0.5 (midpoint).\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" numpy.ndarray: Array of interpolated points with shape (3*N, 6).\n",
|
|
" \"\"\"\n",
|
|
" # Extract coordinates and colors\n",
|
|
" coords = points[:, :3]\n",
|
|
" colors = points[:, 3:]\n",
|
|
" N = coords.shape[0]\n",
|
|
"\n",
|
|
" # Build KDTree for efficient neighbor lookup\n",
|
|
" tree = cKDTree(coords, balanced_tree=False)\n",
|
|
"\n",
|
|
" # Query for 4 nearest neighbors (including self), then exclude self\n",
|
|
" _, indices = tree.query(coords, k=n_c+1, workers=-1)\n",
|
|
" neighbor_indices = indices[:, 1:(n_c+1)] # Shape (N, 3)\n",
|
|
"\n",
|
|
" # Prepare indices for vectorized operations\n",
|
|
" original_indices = np.repeat(np.arange(N), n_c)\n",
|
|
" neighbors_flat = neighbor_indices.ravel()\n",
|
|
"\n",
|
|
" # Gather original and neighbor data\n",
|
|
" original_coords = coords[original_indices]\n",
|
|
" neighbor_coords = coords[neighbors_flat]\n",
|
|
"\n",
|
|
" original_colors = colors[original_indices]\n",
|
|
" neighbor_colors = colors[neighbors_flat]\n",
|
|
"\n",
|
|
" # Interpolate coordinates and colors\n",
|
|
" interpolated_coords = (1 - alpha) * original_coords + alpha * neighbor_coords\n",
|
|
" interpolated_colors = (1 - alpha) * original_colors + alpha * neighbor_colors\n",
|
|
"\n",
|
|
" # Combine into new points array\n",
|
|
" new_points = np.hstack((interpolated_coords, interpolated_colors))\n",
|
|
"\n",
|
|
" combined_points = np.vstack((points, new_points))\n",
|
|
"\n",
|
|
" return combined_points\n",
|
|
"\n",
|
|
"def convolve2d_np(image, kernel, mode='same'):\n",
|
|
" # Flip the kernel for convolution\n",
|
|
" kernel = np.flipud(np.fliplr(kernel))\n",
|
|
" k_h, k_w = kernel.shape\n",
|
|
" i_h, i_w = image.shape\n",
|
|
"\n",
|
|
" # Determine padding\n",
|
|
" if mode == 'same':\n",
|
|
" pad_top = (k_h - 1) // 2\n",
|
|
" pad_bottom = (k_h - 1) - pad_top\n",
|
|
" pad_left = (k_w - 1) // 2\n",
|
|
" pad_right = (k_w - 1) - pad_left\n",
|
|
" padded_image = np.pad(image, ((pad_top, pad_bottom), (pad_left, pad_right)), mode='constant')\n",
|
|
" elif mode == 'valid':\n",
|
|
" padded_image = image\n",
|
|
" elif mode == 'full':\n",
|
|
" padded_image = np.pad(image, ((k_h-1, k_h-1), (k_w-1, k_w-1)), mode='constant')\n",
|
|
" else:\n",
|
|
" raise ValueError(\"Mode must be 'same', 'valid', or 'full'\")\n",
|
|
"\n",
|
|
" # Generate sliding windows\n",
|
|
" windows = np.lib.stride_tricks.sliding_window_view(padded_image, (k_h, k_w))\n",
|
|
"\n",
|
|
" # Perform convolution by summing element-wise multiplication\n",
|
|
" result = np.sum(windows * kernel.reshape(1, 1, k_h, k_w), axis=(-2, -1))\n",
|
|
"\n",
|
|
" return result\n",
|
|
"\n",
|
|
"def blur_image_excluding_black(image, kernel_size=(4,4), threshold=2):\n",
|
|
" # Create mask for non-black pixels (all channels zero)\n",
|
|
" if image.ndim == 3:\n",
|
|
" mask = np.any(image > threshold, axis=-1).astype(float)\n",
|
|
" else:\n",
|
|
" mask = (image > threshold).astype(float)\n",
|
|
"\n",
|
|
" kernel = np.ones(kernel_size)\n",
|
|
"\n",
|
|
" if image.ndim == 3:\n",
|
|
" blurred = np.ones_like(image)*-1\n",
|
|
" for c in range(image.shape[2]):\n",
|
|
" channel = image[:, :, c]\n",
|
|
" masked_channel = channel * mask\n",
|
|
" sum_matrix = convolve2d_np(masked_channel, kernel, mode='same')\n",
|
|
" count_matrix = convolve2d_np(mask, kernel, mode='same')\n",
|
|
" with np.errstate(divide='ignore', invalid='ignore'):\n",
|
|
" mean_matrix = sum_matrix / count_matrix\n",
|
|
" # Where count is zero, use original pixel if non-black, else 0\n",
|
|
" blurred_channel = np.where(count_matrix > 1, mean_matrix, -1)\n",
|
|
" blurred[:, :, c] = blurred_channel\n",
|
|
" else:\n",
|
|
" masked_image = image * mask\n",
|
|
" sum_matrix = convolve2d_np(masked_image, kernel, mode='same')\n",
|
|
" count_matrix = convolve2d_np(mask, kernel, mode='same')\n",
|
|
" with np.errstate(divide='ignore', invalid='ignore'):\n",
|
|
" mean_matrix = sum_matrix / count_matrix\n",
|
|
" blurred = np.where(count_matrix > 0, mean_matrix, masked_image)\n",
|
|
"\n",
|
|
" return blurred"
|
|
],
|
|
"metadata": {
|
|
"id": "PoTzLauTQlHl",
|
|
"cellView": "form"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"\n",
|
|
"from google.colab import files\n",
|
|
"import ipywidgets as widgets\n",
|
|
"from IPython.display import display\n",
|
|
"import PIL\n",
|
|
"from PIL import Image, ImageDraw, ImageFilter\n",
|
|
"\n",
|
|
"#@title #Load image from upload file\n",
|
|
"\n",
|
|
"\n",
|
|
"\n",
|
|
"\n",
|
|
"loadImage = False\n",
|
|
"mask = Image.new('RGBA', (512, 512), (255, 255, 255))\n",
|
|
"initt = None\n",
|
|
"url = None\n",
|
|
"\n",
|
|
"\n",
|
|
"display(\"Load Image:\")\n",
|
|
"Up = files.upload()\n",
|
|
"for k, v in Up.items():\n",
|
|
" initt=Image.open(k)\n",
|
|
"\n",
|
|
"\n",
|
|
"print(initt.size)"
|
|
],
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "4kuU5vFnlfGY"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "6aEoYK7UPbPR",
|
|
"cellView": "form"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"#@title Generate depth map\n",
|
|
"test =True #@param {\"type\":\"boolean\"}\n",
|
|
"# load pipe\n",
|
|
"pipe = pipeline(task=\"depth-estimation\", model=\"depth-anything/Depth-Anything-V2-Large-hf\")\n",
|
|
"\n",
|
|
"# load image\n",
|
|
"if test:\n",
|
|
" url = 'https://raw.githubusercontent.com/chri002/ComfyUI_depthMapOperation/refs/heads/main/assets/start.jpg'\n",
|
|
" image = Image.open(requests.get(url, stream=True).raw)\n",
|
|
" image = image.resize((int(image.size[0]/1.5), int(image.size[1]/1.5)))\n",
|
|
" print(image.size)\n",
|
|
"else:\n",
|
|
" image=initt.convert(\"RGB\")\n",
|
|
"# inference\n",
|
|
"depth_img = pipe(image)[\"depth\"]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"#@title Generate Cloud Points\n",
|
|
"depth =512#@param {\"type\" : \"number\"}\n",
|
|
"\n",
|
|
"quality=2 #@param {\"type\" : \"integer\"}\n",
|
|
"img = torch.from_numpy((np.array(image)))\n",
|
|
"depth_map = torch.from_numpy(np.array(depth_img))\n",
|
|
"points = (from2Dto3D_vectorized_torch(depth_map, img, depth, quality)).cpu().numpy()"
|
|
],
|
|
"metadata": {
|
|
"id": "u9S_ulEFRmiI",
|
|
"cellView": "form"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"#@title Clean Cloud Points\n",
|
|
"\n",
|
|
"k = 20#@param {\"type\" : \"integer\"}\n",
|
|
"distance = 16#@param {\"type\" : \"number\"}\n",
|
|
"\n",
|
|
"points = transform_points(points, translate=(0, 0, 0), rotate=(0, 0, 0), scale=(1,1,2))\n",
|
|
"points = (clean_points(points, k=k, m=distance))\n",
|
|
"points = transform_points(points, translate=(0, 0, 0), rotate=(0, 0, 0), scale=(1,1,0.5))"
|
|
],
|
|
"metadata": {
|
|
"id": "gEjBciYYldPF",
|
|
"cellView": "form"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"#@title Local Rotation\n",
|
|
"rx=0#@param {\"type\" : \"number\"}\n",
|
|
"ry=15#@param {\"type\" : \"number\"}\n",
|
|
"rz=0#@param {\"type\" : \"number\"}\n",
|
|
"\n",
|
|
"mx, Mx, my, My, mz, Mz = points[:,0].min(), points[:,0].max(), points[:,1].min(), points[:,1].max(), points[:,2].min(), points[:,2].max()\n",
|
|
"tx = -(Mx+mx)/2\n",
|
|
"ty = -(My+my)/2\n",
|
|
"tz = -(Mz+mz)/2\n",
|
|
"\n",
|
|
"points = transform_points(points, translate=(tx, ty, tz), rotate=(0, 0, 0), scale=(1,1,1))\n",
|
|
"points = transform_points(points, translate=(0, 0, 0), rotate=(rx/180*m.pi, ry/180*m.pi, rz/180*m.pi), scale=(1,1,1))\n",
|
|
"points = transform_points(points, translate=(-tx, -ty, -tz), rotate=(0, 0, 0), scale=(1,1,1))"
|
|
],
|
|
"metadata": {
|
|
"id": "pV_6QxXTSfF7",
|
|
"cellView": "form"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"#@title Translation\n",
|
|
"tx=0#@param {\"type\" : \"number\"}\n",
|
|
"ty=0#@param {\"type\" : \"number\"}\n",
|
|
"tz=200#@param {\"type\" : \"number\"}\n",
|
|
"\n",
|
|
"points = transform_points(points, translate=(tx, ty, tz), rotate=(0, 0, 0), scale=(1,1,1))"
|
|
],
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "2HcdCFzirZit"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"#@title Scale\n",
|
|
"sx=1#@param {\"type\" : \"number\"}\n",
|
|
"sy=1#@param {\"type\" : \"number\"}\n",
|
|
"sz=1.15#@param {\"type\" : \"number\"}\n",
|
|
"\n",
|
|
"points = transform_points(points, translate=(0,0,0), rotate=(0, 0, 0), scale=(sx,sy,sz))"
|
|
],
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "8Xp1q3DzrONh"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"#@title Points interpolation\n",
|
|
"interpolate = 2 #@param {\"type\":\"number\"}\n",
|
|
"points = torch.from_numpy(interpolate_points(points.detach().cpu().numpy(), alpha=0.5, n_c=interpolate))"
|
|
],
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "uQXuFh6Gr30U"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"#@title Generate image\n",
|
|
"import datetime\n",
|
|
"start = datetime.datetime.now()\n",
|
|
"\n",
|
|
"tend = 0\n",
|
|
"\n",
|
|
"\n",
|
|
"\n",
|
|
"width,height = image.size\n",
|
|
"translation = np.array([-width/2,-height/2,0])\n",
|
|
"scale = np.array([1,1,-1])\n",
|
|
"fov=45 #@param {\"type\" : \"number\"}\n",
|
|
"correction = True #@param{\"type\":\"boolean\"}\n",
|
|
"k_size =3 #@param{\"type\" : \"number\"}\n",
|
|
"points_cop = transform_points(points.copy(), translate=translation, scale=scale)\n",
|
|
"translation = np.array([0,0, -(height/m.sin(m.radians(fov)))])\n",
|
|
"points_cop = transform_points(points_cop, translate=translation)\n",
|
|
"\n",
|
|
"\n",
|
|
"end_img = (render_points_fast(points_cop, img_size=(width,height), color=True, cameraType=1, fov=fov, correction=correction, ksize=(k_size,k_size), thresh=3))\n",
|
|
"\n",
|
|
"background = 125*np.ones((int(image.size[0]/20),int(image.size[1]/20),3))\n",
|
|
"background[1::2, ::2,:]=255\n",
|
|
"background[::2, 1::2,:]=255\n",
|
|
"background= cv2.resize(background, image.size, interpolation = cv2.INTER_NEAREST)\n",
|
|
"\n",
|
|
"mask = end_img[:,:,0]==-1\n",
|
|
"end_img[mask] = background[mask]\n",
|
|
"mask = end_img[:,:,1]==-1\n",
|
|
"end_img[mask] = background[mask]\n",
|
|
"mask = end_img[:,:,2]==-1\n",
|
|
"end_img[mask] = background[mask]\n",
|
|
"\n",
|
|
"tend=datetime.datetime.now()\n",
|
|
"\n",
|
|
"delta=tend-start\n",
|
|
"\n",
|
|
"print(\"vertices: \",points.shape[0],\"\\ntime: \",delta)\n",
|
|
"\n",
|
|
"img_end = Image.fromarray(end_img.astype(np.uint8))\n",
|
|
"\n",
|
|
"images = [image, img_end]\n",
|
|
"widths, heights = zip(*(i.size for i in images))\n",
|
|
"\n",
|
|
"total_width = sum(widths)\n",
|
|
"max_height = max(heights)\n",
|
|
"\n",
|
|
"new_im = Image.new('RGB', (total_width, max_height))\n",
|
|
"\n",
|
|
"x_offset = 0\n",
|
|
"for im in images:\n",
|
|
" new_im.paste(im, (x_offset,0))\n",
|
|
" x_offset += im.size[0]\n",
|
|
"\n",
|
|
"display(new_im)"
|
|
],
|
|
"metadata": {
|
|
"id": "f70devzISk8j",
|
|
"cellView": "form"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"#@title Show Image\n",
|
|
"img_end"
|
|
],
|
|
"metadata": {
|
|
"id": "N56fP1HRrGCD",
|
|
"cellView": "form"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"#@title Export PLY\n",
|
|
"export_PLY(points.detach().cpu().numpy(), name_file=\"model.ply\", multiple_files=False, format_ascii=\"binary_little_endian\")"
|
|
],
|
|
"metadata": {
|
|
"id": "osWtc_Ln0T67",
|
|
"cellView": "form"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"#@title Clean memory\n",
|
|
"clean()"
|
|
],
|
|
"metadata": {
|
|
"id": "B7WfeVFanOqx",
|
|
"cellView": "form"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
}
|
|
]
|
|
} |