Demo that shows how to use TorchGeo to do windowed reading from remote COGs.

torchgeo_rasterdataset_cog.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d17571c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import rasterio\n",
    "from torchgeo.datasets import RasterDataset, stack_samples\n",
    "from torchgeo.samplers import RandomGeoSampler, Units\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "08c2ef04",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"https://maxar-opendata.s3.amazonaws.com/events/Maui-Hawaii-fires-Aug-23/ard/04/122000330002/2023-08-12/10300100EB15FF00-visual.tif\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ec984e36",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(17408, 17408)\n"
     ]
    }
   ],
   "source": [
    "with rasterio.open(url) as f:\n",
    "    print(f.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3a075498",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = RasterDataset(paths=[url])\n",
    "sampler = RandomGeoSampler(ds, size=256, length=32, units=Units.PIXELS)\n",
    "\n",
    "bbox = next(iter(sampler))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c9e29caa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'crs': CRS.from_epsg(32604),\n",
       " 'bbox': BoundingBox(minx=739923.5968245193, maxx=740001.7218245193, miny=2314283.382263966, maxy=2314361.507263966, mint=0.0, maxt=9.223372036854776e+18),\n",
       " 'image': tensor([[[138., 103.,  88.,  ..., 112., 109., 117.],\n",
       "          [162., 133.,  89.,  ..., 117., 109., 116.],\n",
       "          [181., 166., 121.,  ..., 121., 114., 119.],\n",
       "          ...,\n",
       "          [ 58.,  41.,  25.,  ..., 239., 241., 251.],\n",
       "          [ 52.,  38.,  28.,  ..., 219., 225., 235.],\n",
       "          [ 45.,  38.,  30.,  ..., 201., 205., 214.]],\n",
       " \n",
       "         [[130.,  95.,  80.,  ...,  88.,  87.,  97.],\n",
       "          [154., 125.,  82.,  ...,  91.,  87.,  96.],\n",
       "          [174., 159., 114.,  ...,  95.,  90.,  99.],\n",
       "          ...,\n",
       "          [ 66.,  45.,  29.,  ..., 224., 226., 236.],\n",
       "          [ 56.,  42.,  32.,  ..., 204., 208., 220.],\n",
       "          [ 49.,  40.,  32.,  ..., 184., 187., 197.]],\n",
       " \n",
       "         [[107.,  72.,  57.,  ...,  64.,  63.,  73.],\n",
       "          [131., 102.,  56.,  ...,  68.,  63.,  71.],\n",
       "          [148., 133.,  88.,  ...,  72.,  66.,  74.],\n",
       "          ...,\n",
       "          [ 45.,  28.,  12.,  ..., 203., 205., 215.],\n",
       "          [ 39.,  25.,  15.,  ..., 183., 188., 199.],\n",
       "          [ 32.,  26.,  18.,  ..., 164., 167., 177.]]])}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds[bbox]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0b4e68a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 256, 256])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds[bbox][\"image\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1f7af06b",
   "metadata": {},
   "outputs": [],
   "source": [
    "dl = DataLoader(ds, batch_size=4, sampler=sampler, num_workers=6, collate_fn=stack_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9a4b6c22",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 3, 256, 256])\n",
      "torch.Size([4, 3, 256, 256])\n",
      "torch.Size([4, 3, 256, 256])\n",
      "torch.Size([4, 3, 256, 256])\n",
      "torch.Size([4, 3, 256, 256])\n",
      "torch.Size([4, 3, 256, 256])\n",
      "torch.Size([4, 3, 256, 256])\n",
      "torch.Size([4, 3, 256, 256])\n"
     ]
    }
   ],
   "source": [
    "for batch in dl:\n",
    "    print(batch[\"image\"].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac6865a3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}