Update README.md

README.md

@@ -12,12 +12,12 @@ SegFormer model with a MiT-b2 backbone fine-tuned on Coralscapes at resolution 1
 
 ### Model Description
 
-Training is conducted following \cite{xie2021segformer}, using a batch size of 8 for 265 epochs, 
-using the AdamW optimizer with an initial learning rate of 6e-5 (multiplied by 10 when using LoRA), 
-weight decay of 1e-2 and polynomial learning rate scheduler with a power of 1. 
+Training is conducted following the Segformer original [implementation](https://proceedings.neurips.cc/paper_files/paper/2021/file/64f1f27bf1b4ec22924fd0acb550c235-Paper.pdf), using a batch size of 8 for 265 epochs, 
+using the AdamW optimizer with an initial learning rate of 6e-5, weight decay of 1e-2 and polynomial learning rate scheduler with a power of 1. 
 During training, images are randomly scaled within a range of 1 and 2, flipped horizontally with a 0.5 probability and randomly cropped to 1024×1024 pixels. 
 Input images are normalized using the ImageNet mean and standard deviation. For evaluation, a non-overlapping sliding window strategy is employed, 
-using a window size of 1024x1024 and a stride of 1024. 
+using a window size of 1024x1024. 
+<!-- TODO - We used a stride of 1024 but in the demo it is variable. Should we move this entire section to training below? -->
 
 - **Developed by:** [More Information Needed]
 - **Funded by [optional]:** [More Information Needed]
@@ -30,9 +30,9 @@ using a window size of 1024x1024 and a stride of 1024.
 
 <!-- Provide the basic links for the model. -->
 
-- **Repository:** [coralscapesScripts](https://github.com/ViktorDomazetoski/coralscapesScripts/)
+- **Repository:** [coralscapesScripts](https://github.com/eceo-epfl/coralscapesScripts/)
 - **Paper [optional]:** [More Information Needed]
-- **Demo [optional]:** [More Information Needed]
+- **Demo** [Hugging Face Spaces](https://huggingface.co/spaces/EPFL-ECEO/coralscapes_demo):
 
 ## Uses
 
@@ -70,32 +70,113 @@ Users (both direct and downstream) should be made aware of the risks, biases and
 
 ## How to Get Started with the Model
 
-Here is how to use this model to segment an image of the Coralscapes dataset
+The simplest way to use this model to segment an image of the Coralscapes dataset is as follows:
 
 ```python
 from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
 from PIL import Image
 from datasets import load_dataset
-import numpy as np
 
 # Load an image from the coralscapes dataset or load your own image 
 dataset = load_dataset("EPFL-ECEO/coralscapes") 
 image = dataset["test"][42]["image"]
-label = dataset["test"][42]["label"]
-
-id2label = {"0": "unlabeled", "1": "seagrass", "2": "trash", "3": "other coral dead", "4": "other coral bleached", "5": "sand", "6": "other coral alive", "7": "human", "8": "transect tools", "9": "fish", "10": "algae covered substrate", "11": "other animal", "12": "unknown hard substrate", "13": "background", "14": "dark", "15": "transect line", "16": "massive/meandering bleached", "17": "massive/meandering alive", "18": "rubble", "19": "branching bleached", "20": "branching dead", "21": "millepora", "22": "branching alive", "23": "massive/meandering dead", "24": "clam", "25": "acropora alive", "26": "sea cucumber", "27": "turbinaria", "28": "table acropora alive", "29": "sponge", "30": "anemone", "31": "pocillopora alive", "32": "table acropora dead", "33": "meandering bleached", "34": "stylophora alive", "35": "sea urchin", "36": "meandering alive", "37": "meandering dead", "38": "crown of thorn", "39": "dead clam"}
-label2color = {"unlabeled":[255, 255, 255], "human": [255, 0, 0], "background": [29, 162, 216], "fish": [255, 255, 0], "sand": [194, 178, 128], "rubble": [161, 153, 128], "unknown hard substrate": [125, 125, 125], "algae covered substrate": [125, 163, 125], "dark": [31, 31, 31], "branching bleached": [252, 231, 240], "branching dead": [123, 50, 86], "branching alive": [226, 91, 157], "stylophora alive": [255, 111, 194], "pocillopora alive": [255, 146, 150], "acropora alive": [236, 128, 255], "table acropora alive": [189, 119, 255], "table acropora dead": [85, 53, 116], "millepora": [244, 150, 115], "turbinaria": [228, 255, 119], "other coral bleached": [250, 224, 225], "other coral dead": [114, 60, 61], "other coral alive": [224, 118, 119], "massive/meandering alive": [236, 150, 21], "massive/meandering dead": [134, 86, 18], "massive/meandering bleached": [255, 248, 228], "meandering alive": [230, 193, 0], "meandering dead": [119, 100, 14], "meandering bleached": [251, 243, 216], "transect line": [0, 255, 0], "transect tools": [8, 205, 12], "sea urchin": [0, 142, 255], "sea cucumber": [0, 231, 255], "anemone": [0, 255, 189], "sponge": [240, 80, 80], "clam": [189, 255, 234], "other animal": [0, 255, 255], "trash": [255, 0, 134], "seagrass": [125, 222, 125], "crown of thorn": [179, 245, 234], "dead clam": [89, 155, 134]}
-id2color = {int(id): label2color[label] for id, label in id2label.items()}
 
 preprocessor = SegformerImageProcessor.from_pretrained("EPFL-ECEO/segformer-b2-finetuned-coralscapes-1024-1024")
 model = SegformerForSemanticSegmentation.from_pretrained("EPFL-ECEO/segformer-b2-finetuned-coralscapes-1024-1024")
 
-inputs = preprocessor(image.resize((1024, 512)), return_tensors = "pt")
+inputs = preprocessor(image, return_tensors = "pt")
 outputs = model(**inputs)
 outputs = preprocessor.post_process_semantic_segmentation(outputs, target_sizes=[(image.size[1], image.size[0])])
 label_pred = outputs[0].cpu().numpy()
 ```
 
+While using the above approach should still work for images of different sizes and scales, for images that are not close to the training size of the model (1024x1024), 
+we recommend using the following approach using a sliding window to achieve better results:
+
+```python
+import torch 
+import torch.nn.functional as F
+from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
+from PIL import Image
+from datasets import load_dataset
+import numpy as np
+
+def resize_image(image, target_size=1024):
+    """
+    Used to resize the image such that the smaller side equals 1024
+    """
+    h_img, w_img = image.size
+    if h_img < w_img:
+        new_h, new_w = target_size, int(w_img * (target_size / h_img))
+    else:
+        new_h, new_w  = int(h_img * (target_size / w_img)), target_size
+    resized_img = image.resize((new_h, new_w))
+    return resized_img
+
+def segment_image(image, preprocessor, model, crop_size = (1024, 1024), num_classes = 40, transform=None):
+    """
+    Finds an optimal stride based on the image size and aspect ratio to create
+    overlapping sliding windows of size 1024x1024 which are then fed into the model.  
+    """ 
+    h_crop, w_crop = crop_size
+    
+    img = torch.Tensor(np.array(resize_image(image, target_size=1024)).transpose(2, 0, 1)).unsqueeze(0)
+    batch_size, _, h_img, w_img = img.size()
+    
+    if transform:
+        img = torch.Tensor(transform(image = img.numpy())["image"]).to(device)    
+        
+    h_grids = int(np.round(3/2*h_img/h_crop)) if h_img > h_crop else 1
+    w_grids = int(np.round(3/2*w_img/w_crop)) if w_img > w_crop else 1
+    
+    h_stride = int((h_img - h_crop + h_grids -1)/(h_grids -1)) if h_grids > 1 else h_crop
+    w_stride = int((w_img - w_crop + w_grids -1)/(w_grids -1)) if w_grids > 1 else w_crop
+    
+    preds = img.new_zeros((batch_size, num_classes, h_img, w_img))
+    count_mat = img.new_zeros((batch_size, 1, h_img, w_img))
+    
+    for h_idx in range(h_grids):
+        for w_idx in range(w_grids):
+          y1 = h_idx * h_stride
+          x1 = w_idx * w_stride
+          y2 = min(y1 + h_crop, h_img)
+          x2 = min(x1 + w_crop, w_img)
+          y1 = max(y2 - h_crop, 0)
+          x1 = max(x2 - w_crop, 0)
+          crop_img = img[:, :, y1:y2, x1:x2]
+          with torch.no_grad():
+            if(preprocessor):
+                inputs = preprocessor(crop_img, return_tensors = "pt")
+                inputs["pixel_values"] = inputs["pixel_values"].to(device)
+            else:
+                inputs = crop_img.to(device)
+            outputs = model(**inputs)
+
+          resized_logits = F.interpolate(
+              outputs.logits[0].unsqueeze(dim=0), size=crop_img.shape[-2:], mode="bilinear", align_corners=False
+          )
+          preds += F.pad(resized_logits,
+                          (int(x1), int(preds.shape[3] - x2), int(y1),
+                          int(preds.shape[2] - y2)))
+          count_mat[:, :, y1:y2, x1:x2] += 1
+    
+    assert (count_mat == 0).sum() == 0
+    preds = preds / count_mat
+    preds = preds.argmax(dim=1)
+    preds = F.interpolate(preds.unsqueeze(0).type(torch.uint8), size=image.size[::-1], mode='nearest')
+    label_pred = preds.squeeze().cpu().numpy()
+    return label_pred
+
+# Load an image from the coralscapes dataset or load your own image 
+dataset = load_dataset("EPFL-ECEO/coralscapes") 
+image = dataset["test"][42]["image"]
+
+preprocessor = SegformerImageProcessor.from_pretrained("EPFL-ECEO/segformer-b2-finetuned-coralscapes-1024-1024")
+model = SegformerForSemanticSegmentation.from_pretrained("EPFL-ECEO/segformer-b2-finetuned-coralscapes-1024-1024")
+
+label_pred = segment_image(image, preprocessor, model)
+```
+
 ## Training Details
 
 ### Training Data