update readme

README.md

@@ -1,5 +1,5 @@
 ---
-title: GRowSeg demo
+title: Vineyard Row Segmentation
 emoji: 🍇
 colorFrom: indigo
 colorTo: gray
@@ -14,6 +14,8 @@ models:
 - links-ads/gaia-growseg
 datasets:
 - links-ads/gaia-vineyard-uav-dataset
+preload_from_hub:
+- links-ads/gaia-growseg model.safetensors
 tags:
 - agriculture
 - viticulture

lib/utils.py

@@ -80,36 +80,52 @@ def pad(img, pad, order='CHW'):
     return padded_img
 
 
-def extract_patches(img, patch_size=512, stride=256, order='CHW', only_return_idx=True):
+def extract_patches(img, patch_size=512, stride=256, order='CHW', only_return_idx=True, include_last=True):
     """Extract patches from an image, in the format (h_start, h_end, w_start, w_end)"""
     assert order in ['HWC', 'CHW'], f"Got unknown order '{order}', expected one of ['HWC','CHW']"
-
+    assert len(img.shape) == 3, f"Got image with {len(img.shape)} dimensions, expected 3 dimensions (C,H,W) or (H,W,C)"
+    assert img.shape[0] >= patch_size, f"Got image with height {img.shape[0]}, expected at least {patch_size}. Maybe apply padding first?"
+    assert img.shape[1] >= patch_size, f"Got image with width {img.shape[1]}, expected at least {patch_size}. Maybe apply padding first?"
+    
+    # Get image height and width
     if order == 'HWC':
         H, W = img.shape[:2]
     else:
         H, W = img.shape[1:]
 
-    # compute the number of patches
-    n_patches = ((H - patch_size) // stride + 1) * ((W - patch_size) // stride + 1)
+    # Compute the number of "proper" patches in each dimension
+    n_patches_H = (H - patch_size) // stride + 1
+    n_patches_W = (W - patch_size) // stride + 1
 
-    # extract patches
-    patches = []
+    # Extract patches indices
     patches_idx = []
-    for i in range(0, H-patch_size+1, stride):
-        for j in range(0, W-patch_size+1, stride):
-            
-            patches_idx.append((i, i+patch_size, j, j+patch_size))
-
-            if not only_return_idx:
-                if order == 'HWC':
-                    patch = img[i:i+patch_size, j:j+patch_size, :]
-                else:
-                    patch = img[:, i:i+patch_size, j:j+patch_size]
-                patches.append(patch)
-            
+    for i in range(n_patches_H):    # iterate over height
+        for j in range(n_patches_W):    # iterate over width
+
+            # Get the current patch indices
+            patches_idx.append((i*stride, i*stride+patch_size, j*stride, j*stride+patch_size))  # (top, bottom, left, right)
+
+            # Include leftmost and lowermost patch if needed
+            if include_last:
+                if j == n_patches_W-1 and j*stride+patch_size < W:
+                    patches_idx.append((i*stride, i*stride+patch_size, W-patch_size, W))
+                if i == n_patches_H-1 and i*stride+patch_size < H:
+                    patches_idx.append((H-patch_size, H, j*stride, j*stride+patch_size))
+                if i == n_patches_H-1 and j == n_patches_W-1 and i*stride+patch_size < H and j*stride+patch_size < W:
+                    patches_idx.append((H-patch_size, H, W-patch_size, W))
+
     if only_return_idx:
         return patches_idx
-    return patches, patches_idx
+    else:
+        # Extract patches
+        patches = []
+        for t,b,l,r in patches_idx:
+            if order == 'HWC':
+                patch = img[t:b, l:r, :]
+            else:
+                patch = img[:, t:b, l:r]
+            patches.append(patch)
+        return patches, patches_idx
 
 
 def segment_batch(batch, model):