update datamodule and tutorial

Author: neerajvk10

claymodel/datamodule.py

@@ -244,7 +244,7 @@ def __init__(  # noqa: PLR0913
         self.batch_size = batch_size
         self.num_workers = num_workers
         self.prefetch_factor = prefetch_factor
-        self.split_ratio = 0.8
+        self.split_ratio = 0.7
 
     def setup(self, stage: Literal["fit", "predict"] | None = None) -> None:
         # Get list of GeoTIFF filepaths from s3 bucket or data/ folder
@@ -258,13 +258,24 @@ def setup(self, stage: Literal["fit", "predict"] | None = None) -> None:
         print(f"Total number of chips: {len(chips_path)}")
 
         if stage == "fit":
-            trn_paths, val_paths = train_test_split(
-                chips_path,
-                test_size=(1 - self.split_ratio),
-                stratify=chips_platform,
-                shuffle=True,
-            )
-
+        # Check how many unique platforms we have
+            unique_platforms = set(chips_platform)
+
+            if len(unique_platforms) > 1:
+                # Use stratification if multiple platforms exist
+                trn_paths, val_paths = train_test_split(
+                    chips_path,
+                    test_size=(1 - self.split_ratio),
+                    stratify=chips_platform,
+                    shuffle=True,
+                )
+            else:
+                # Disable stratification if only one platform (e.g. NAIP)
+                trn_paths, val_paths = train_test_split(
+                    chips_path,
+                    test_size=(1 - self.split_ratio),
+                    shuffle=True,
+                )
             self.trn_ds = EODataset(
                 chips_path=trn_paths,
                 size=self.size,

docs/tutorials/wall-to-wall.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 45,
+   "execution_count": 1,
    "id": "b958c15a",
    "metadata": {},
    "outputs": [],
@@ -13,10 +13,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 46,
+   "execution_count": 3,
    "id": "648742f7",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/opt/anaconda3/envs/claymodel/lib/python3.11/site-packages/pyproj/network.py:59: UserWarning: pyproj unable to set PROJ database path.\n",
+      "  _set_context_ca_bundle_path(ca_bundle_path)\n"
+     ]
+    }
+   ],
    "source": [
     "\n",
     "#import os\n",
@@ -30,10 +39,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 47,
+   "execution_count": 4,
    "id": "de0b279c",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/opt/anaconda3/envs/claymodel/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
    "source": [
     "import math\n",
     "import geopandas as gpd\n",