{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":91498,"databundleVersionId":11655853,"sourceType":"competition"}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"! uv pip install --system jaxtyping rerun-sdk[notebook]\n\n! uv pip install --system https://github.com/facebookresearch/vggt.git","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from dataclasses import dataclass\n\nimport torch\nfrom vggt.models.vggt import VGGT\nfrom vggt.utils.geometry import unproject_depth_map_to_point_map\nfrom vggt.utils.load_fn import load_and_preprocess_images\nfrom vggt.utils.pose_enc import pose_encoding_to_extri_intri\n\n\n@dataclass\nclass VGGTResult:\n    inputs: torch.Tensor  # [num_images, 3, H, W]\n    width: int\n    height: int\n    extrinsic: torch.Tensor  # [num_images, 3, 4]\n    intrinsic: torch.Tensor  # [num_images, 3, 3]\n    depth_map: torch.Tensor  # [num_images, H, W, 1]\n    depth_conf: torch.Tensor  # [num_images, H, W]\n    point_map: torch.Tensor  # [num_images, H, W, 3]\n    point_conf: torch.Tensor  # [num_images, H, W]\n    point_map_by_unprojection: torch.Tensor  # [num_images, H, W, 3]\n\n    @property\n    def images(self) -> torch.Tensor:\n        \"\"\"Convert the inputs to images for visualization.\n\n        Returns:\n            torch.Tensor: The images in uint8 format. [num_images, H, W, 3].\n        \"\"\"\n        return (self.inputs.permute(0, 2, 3, 1) * 255).clamp(0, 255).to(torch.uint8)\n\n\nclass VGGTRunner:\n    \"\"\"VGGT Runner\n\n    Attributes:\n        model_path (str): The path to the VGGT model.\n        dtype (torch.dtype | None): The dtype of the model.\n        device (torch.device | str | None): The device of the model.\n    \"\"\"\n\n    def __init__(\n        self,\n        model_path: str = \"facebook/VGGT-1B\",\n        dtype: torch.dtype | None = None,\n        device: torch.device | str | None = None,\n    ):\n        \"\"\"Initialize the VGGT Runner.\n\n        Args:\n            model_path (str): The path to the VGGT model.\n            dtype (torch.dtype | None, optional): The dtype for autocasting. Defaults to None.\n            device (torch.device | str | None, optional): The device of the model. Defaults to None.\n        \"\"\"\n        self.dtype = dtype if dtype is not None else torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16\n        self.device = torch.device(device) if device is not None else torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        self.model = VGGT.from_pretrained(model_path).eval().to(self.device)\n\n    def __call__(self, image_paths: list[str]) -> VGGTResult:\n        \"\"\"Inference the VGGT model.\n\n        Args:\n            image_paths (list[str]): The paths to the images.\n\n        Returns:\n            VGGTResult: The result of the inference.\n        \"\"\"\n        inputs = load_and_preprocess_images(image_paths).to(self.device).unsqueeze(0)\n        with torch.inference_mode():\n            with torch.amp.autocast(str(self.device), dtype=self.dtype):\n                aggregated_tokens_list, ps_idx = self.model.aggregator(inputs)\n            pose_enc = self.model.camera_head(aggregated_tokens_list)[-1]\n            extrinsic, intrinsic = pose_encoding_to_extri_intri(pose_enc, inputs.shape[-2:])\n            depth_map, depth_conf = self.model.depth_head(aggregated_tokens_list, inputs, ps_idx)\n            point_map, point_conf = self.model.point_head(aggregated_tokens_list, inputs, ps_idx)\n            point_map_by_unprojection = unproject_depth_map_to_point_map(depth_map.squeeze(0), extrinsic.squeeze(0), intrinsic.squeeze(0))\n        return VGGTResult(\n            inputs=inputs.cpu().squeeze(0),\n            width=inputs.shape[3],\n            height=inputs.shape[2],\n            extrinsic=extrinsic.cpu().squeeze(0),\n            intrinsic=intrinsic.cpu().squeeze(0),\n            depth_map=depth_map.cpu().squeeze(0),\n            depth_conf=depth_conf.cpu().squeeze(0),\n            point_map=point_map.cpu().squeeze(0),\n            point_conf=point_conf.cpu().squeeze(0),\n            point_map_by_unprojection=torch.from_numpy(point_map_by_unprojection),\n        )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T00:50:10.694361Z","iopub.execute_input":"2025-09-24T00:50:10.694642Z","iopub.status.idle":"2025-09-24T00:50:23.942622Z","shell.execute_reply.started":"2025-09-24T00:50:10.694608Z","shell.execute_reply":"2025-09-24T00:50:23.942043Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from glob import glob\n\nvggt = VGGTRunner(device=\"cuda\")\n\nimage_paths = sorted(glob(\"/kaggle/input/image-matching-challenge-2025/train/imc2024_dioscuri_baalshamin/baalshamin_1*.png\"))\nresult = vggt(image_paths)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T00:59:07.325987Z","iopub.execute_input":"2025-09-24T00:59:07.326747Z","iopub.status.idle":"2025-09-24T00:59:37.411578Z","shell.execute_reply.started":"2025-09-24T00:59:07.326723Z","shell.execute_reply":"2025-09-24T00:59:37.410960Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport rerun as rr\nimport rerun.blueprint as rrb\nfrom tqdm import tqdm\n\n\ndef draw_pose(transform: np.ndarray, name: str, static: bool = False) -> None:\n    rr.log(\n        name,\n        rr.Arrows3D(\n            origins=[0, 0, 0],\n            vectors=[[0.03, 0, 0], [0, 0.03, 0], [0, 0, 0.03]],\n            colors=[[255, 0, 0], [0, 255, 0], [0, 0, 255]],\n            radii=[0.001, 0.001, 0.001],\n        ),\n        static=static,\n    )\n\n    rr.log(\n        name,\n        rr.Transform3D(\n            translation=transform[:3, 3],\n            mat3x3=transform[:3, :3],\n        ),\n        static=static,\n    )\n\n\nrr.init(\"vggt_inference\")\nrr.send_blueprint(\n    rrb.Blueprint(\n        rrb.Horizontal(\n            rrb.Spatial3DView(origin=\"body/pose\", contents=\"body/**\"),\n            rrb.Vertical(\n                rrb.Spatial2DView(origin=\"body/cam/image\"),\n                rrb.Spatial2DView(origin=\"body/cam/depth_map\"),\n            ),\n        )\n    )\n)\n\nfor i in tqdm(range(len(image_paths))):\n    rr.set_time_sequence(\"frame\", i)\n    world2cam = np.eye(4)\n    world2cam[:3, :4] = result.extrinsic[i]\n    cam2world = np.linalg.inv(world2cam)\n    rr.log(\n        \"body/cam\",\n        rr.Pinhole(\n            image_from_camera=result.intrinsic[i],\n            width=result.width,\n            height=result.height,\n            image_plane_distance=0.02,\n        ),\n    )\n    rr.log(\n        \"body/cam\",\n        rr.Transform3D(\n            translation=cam2world[:3, 3],\n            mat3x3=cam2world[:3, :3],\n        ),\n    )\n    draw_pose(cam2world, f\"body/pose{i}\", static=True)\n    draw_pose(cam2world, \"body/pose\")\n\n    # Filter points based on confidence threshold\n    conf = result.depth_conf[i]\n    conf_thres = np.percentile(conf, 90)\n    keep_mask = conf > conf_thres\n    depth_maps = result.depth_map[i]\n    depth_maps[~keep_mask] = 0\n\n    points = result.point_map_by_unprojection[i].reshape(-1, 3)\n    keep_mask = keep_mask.reshape(-1)\n    colors = result.images[i].reshape(-1, 3)\n\n    points = points[keep_mask]\n    colors = colors[keep_mask]\n\n    rr.log(f\"body/points{i}\", rr.Points3D(points, colors=colors, radii=0.0003), static=True)\n    rr.log(\"body/cam/image\", rr.Image(result.images[i]))\n    rr.log(\"body/cam/depth_map\", rr.DepthImage(depth_maps))\n\nrr.save(\"output.rrd\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T00:59:49.594906Z","iopub.execute_input":"2025-09-24T00:59:49.595641Z","iopub.status.idle":"2025-09-24T00:59:52.037954Z","shell.execute_reply.started":"2025-09-24T00:59:49.595610Z","shell.execute_reply":"2025-09-24T00:59:52.037402Z"}},"outputs":[],"execution_count":null}]}