%load_ext autoreload
%autoreload 2

!rm -rf data

import os
import glob
import random
import shutil

max_numfiles = 10
train_fraction = 0.8

datarootdir = "/mounts/Datasets3/2024-2025-ChallengePlankton/train"

if "TMPDIR" in os.environ:
    # On the DCE, $TMPDIR refer to a temporary local storage on the compute node
    tmpdir = os.environ["TMPDIR"]
    os.symlink(tmpdir, "./data")
else:
    tmpdir = "./data"

if os.path.exists(tmpdir):
    shutil.rmtree(tmpdir)

os.makedirs(tmpdir)

# Locate all the xxxxx_scan.png.ppm from the datarootdir with the labels xxxxx_mask.png.ppm
# split it randomly in 80% train and 20% test
# Only consider up to max_numfiles files
# Copy these data in their respective folders into the tmpdir

scan_files = glob.glob(os.path.join(datarootdir, "*_scan.png.ppm"))
# Compute the mask filenames from the scan files
mask_files = [f.replace("_scan.png.ppm", "_mask.png.ppm") for f in scan_files]

# Random shuffle the files for the split. Apply the same random to the masks to keep
# both lists aligned

idx = list(range(len(scan_files)))[:max_numfiles]
random.shuffle(idx)

train_idx = idx[: int(len(idx) * train_fraction)]
test_idx = idx[int(len(idx) * train_fraction) :]

train_scan_files = [scan_files[i] for i in train_idx]
train_mask_files = [mask_files[i] for i in train_idx]
test_scan_files = [scan_files[i] for i in test_idx]
test_mask_files = [mask_files[i] for i in test_idx]


def copy_files(files, target_dir):
    """
    Copy the files to the target directory.
    """
    if not os.path.exists(target_dir):
        os.makedirs(target_dir)
    for f in files:
        os.system(f"cp {f} {target_dir}")


# Copy the files to the tmpdir
copy_files(train_scan_files, os.path.join(tmpdir, "train"))
copy_files(train_mask_files, os.path.join(tmpdir, "train"))

copy_files(test_scan_files, os.path.join(tmpdir, "test"))
copy_files(test_mask_files, os.path.join(tmpdir, "test"))

# Copy the taxa file as well
os.system(f"cp {datarootdir}/taxa.csv {os.path.join(tmpdir, 'train')}")
os.system(f"cp {datarootdir}/taxa.csv {os.path.join(tmpdir, 'test')}")

0

!ls data/train

rg20090204_mask.png.ppm  rg20090617_mask.png.ppm  rg20091028_mask.png.ppm
rg20090204_scan.png.ppm  rg20090617_scan.png.ppm  rg20091028_scan.png.ppm
rg20090318_mask.png.ppm  rg20090715_mask.png.ppm  rg20091125_mask.png.ppm
rg20090318_scan.png.ppm  rg20090715_scan.png.ppm  rg20091125_scan.png.ppm
rg20090610_mask.png.ppm  rg20090902_mask.png.ppm  taxa.csv
rg20090610_scan.png.ppm  rg20090902_scan.png.ppm

!for f in data/train/*_scan.png.ppm; do file $f; done

data/train/rg20090204_scan.png.ppm: Netpbm image data, size = 22817 x 14569, rawbits, greymap
data/train/rg20090318_scan.png.ppm: Netpbm image data, size = 22797 x 14549, rawbits, greymap
data/train/rg20090610_scan.png.ppm: Netpbm image data, size = 22717 x 14459, rawbits, greymap
data/train/rg20090617_scan.png.ppm: Netpbm image data, size = 22817 x 14529, rawbits, greymap
data/train/rg20090715_scan.png.ppm: Netpbm image data, size = 22807 x 14549, rawbits, greymap
data/train/rg20090902_scan.png.ppm: Netpbm image data, size = 22807 x 14319, rawbits, greymap
data/train/rg20091028_scan.png.ppm: Netpbm image data, size = 22787 x 14379, rawbits, greymap
data/train/rg20091125_scan.png.ppm: Netpbm image data, size = 22817 x 14519, rawbits, greymap

!cat data/train/taxa.csv

label,living,label_nb
artefact,False,1
badfocus<artefact,False,2
bubble,False,3
detritus,False,4
fiber<detritus,False,5
t001,False,6
t003,False,7
Abylopsis tetragona,True,8
Acartiidae,True,9
Actiniaria,True,10
Actinopterygii,True,11
Aglaura,True,12
Amphipoda,True,13
Annelida,True,14
Atlanta,True,15
Aulacantha,True,16
Bivalvia<Mollusca,True,17
Calanidae,True,18
Calanoida,True,19
Calocalanus pavo,True,20
Candaciidae,True,21
Cavolinia inflexa,True,22
Centropagidae,True,23
Chaetognatha,True,24
Chelophyes appendiculata,True,25
Cliidae,True,26
Collodaria,True,27
Copilia,True,28
Corycaeidae,True,29
Creseidae,True,30
Creseis acicula,True,31
Ctenophora<Metazoa,True,32
Cymbulia peroni,True,33
Doliolida,True,34
Echinodermata,True,35
Eucalanidae,True,36
Euchaetidae,True,37
Eumalacostraca,True,38
Evadne,True,39
Fritillariidae,True,40
Gammaridea,True,41
Globigerinidae,True,42
Gymnosomata,True,43
Haloptilus,True,44
Harpacticoida,True,45
Heterorhabdidae,True,46
Hydrozoa,True,47
Hyperiidea,True,48
Insecta,True,49
Lensia subtilis,True,50
Limacinidae,True,51
Liriope<Geryoniidae,True,52
Metridinidae,True,53
Mollusca,True,54
Obelia,True,55
Oikopleuridae,True,56
Oithonidae,True,57
Oncaeidae,True,58
Orbulina,True,59
Ostracoda,True,60
Penilia avirostris,True,61
Physonectae,True,62
Podon,True,63
Pontellidae,True,64
Pontellina plumata,True,65
Pyrosomatida,True,66
Rhincalanidae,True,67
Rhizaria,True,68
Rhopalonema velatum,True,69
Salpida,True,70
Sapphirinidae,True,71
Solmundella bitentaculata,True,72
Temoridae,True,73
bract<Abylidae,True,74
bract<Abylopsis tetragona,True,75
bract<Diphyidae,True,76
calyptopsis<Euphausiacea,True,77
chain<Salpida,True,78
cirrus,True,79
colony<Phaeodaria,True,80
cypris,True,81
damaged<Aulacantha,True,82
division,True,83
egg<Actinopterygii,True,84
egg<Mollusca,True,85
egg<other,True,86
ephyra<Scyphozoa,True,87
eudoxie<Abylopsis tetragona,True,88
eudoxie<Bassia bassensis,True,89
eudoxie<Diphyidae,True,90
gonophore<Abylopsis tetragona,True,91
gonophore<Bassia bassensis,True,92
gonophore<Diphyidae,True,93
head<Chaetognatha,True,94
juvenile<Salpida,True,95
larvae<Annelida,True,96
larvae<Porcellanidae,True,97
larvae<Squillidae,True,98
like<Collodaria,True,99
like<Laomediidae,True,100
megalopa,True,101
multiple<Copepoda,True,102
multiple<other,True,103
nauplii<Cirripedia,True,104
nauplii<Crustacea,True,105
nectophore<Abylopsis tetragona,True,106
nectophore<Bassia bassensis,True,107
nectophore<Diphyidae,True,108
nectophore<Hippopodiidae,True,109
nectophore<Physonectae,True,110
nucleus<Salpidae,True,111
othertocheck,True,112
part<Annelida,True,113
part<Cnidaria,True,114
part<Crustacea,True,115
part<Ctenophora,True,116
part<Mollusca,True,117
part<Siphonophorae,True,118
phyllosoma,True,119
pluteus<Echinoidea,True,120
pluteus<Ophiuroidea,True,121
protozoea<Mysida,True,122
protozoea<Penaeidae,True,123
protozoea<Sergestidae,True,124
scale,True,125
seaweed,True,126
siphonula,True,127
t002,True,128
tail<Appendicularia,True,129
tail<Chaetognatha,True,130
trunk<Appendicularia,True,131
wing,True,132
zoea<Brachyura,True,133
zoea<Galatheidae,True,134

# Standard imports
import pathlib

# External imports
import albumentations as A
import numpy as np
import matplotlib.pyplot as plt
import torch
import yaml

# Local imports
import wandb
import planktoseg
from planktoseg import data
from planktoseg import models
from planktoseg import optim
from planktoseg import utils
from planktoseg import metrics
from planktoseg import main

use_cuda = torch.cuda.is_available()
device = torch.device("cuda") if use_cuda else torch.device("cpu")

logdir = pathlib.Path("./logs")
if not logdir.exists():
    logdir.mkdir(exist_ok=True)

/usr/users/dce-admin/fix/GIT/2024_ml4oceans/venv/lib/python3.10/site-packages/requests/__init__.py:86: RequestsDependencyWarning: Unable to find acceptable character detection dependency (chardet or charset_normalizer).
  warnings.warn(

transform = A.pytorch.ToTensorV2()

dataset = data.PlanktonDataset(root="./data/train", patch_size=(512, 512), patch_stride=(128, 128), task=data.SegmentationTask.LIVING_NONLIVING)
dataset = data.WrappedDataset(dataset, transform)

100%|██████████| 8/8 [00:00<00:00, 7584.64it/s]
100%|██████████| 8/8 [00:00<00:00, 7584.64it/s]

idx = np.random.randint(len(dataset)) 

img, mask = dataset[idx]

print(f"The img is of type {type(img)} and of shape {img.shape}")
print(f"The mask is of type {type(mask)} and of shape {mask.shape}")

plt.subplot(1, 2, 1)
plt.imshow(img.squeeze(), cmap="gray")
plt.title("Zooscan image")
plt.axis("off")

plt.subplot(1, 2, 2)
plt.imshow(mask.squeeze(), interpolation="none", cmap="tab20c")
plt.title("Mask")
plt.axis("off")

The img is of type <class 'torch.Tensor'> and of shape torch.Size([1, 512, 512])
The mask is of type <class 'torch.Tensor'> and of shape torch.Size([512, 512])

(np.float64(-0.5), np.float64(511.5), np.float64(511.5), np.float64(-0.5))

idx = np.random.randint(len(dataset)) 
idx = 15000

# Sample the dataset without all the transforms
# to see what the image and mask originally look like
original_transform = A.pytorch.ToTensorV2()
dataset.transform = original_transform
orig_img, orig_mask = dataset[idx]

##########################################################################################
# TODO: Feel free to change the code below 
# Tune the augmented transform by prepending your choosen transform before the 
# conversion to pytorch tensor
# Fill free to evaluate the cell as you add transforms to see their effect
transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.RandomRotate90(p=0.5),
    A.MaskDropout((1, 1), image_fill_value=255, p=1), 
    A.Blur(),
    A.pytorch.ToTensorV2()
])
##########################################################################################

# And now sample the exact same image/mask 
dataset.transform = transform
aug_img, aug_mask = dataset[idx]

plt.subplot(2, 2, 1)
plt.imshow(orig_img.squeeze(), cmap="gray", clim=(0, 255))
plt.title("Zooscan original")
plt.axis("off")

plt.subplot(2, 2, 2)
plt.imshow(orig_mask.squeeze(), interpolation="none", cmap="tab20c")
plt.title("Original Mask")
plt.axis("off")

plt.subplot(2, 2, 3)
plt.imshow(aug_img.squeeze(), cmap="gray", clim=(0, 255))
plt.title("Zooscan augmented image")
plt.axis("off")

plt.subplot(2, 2, 4)
plt.imshow(aug_mask.squeeze(), interpolation="none", cmap="tab20c")
plt.title("Augmented mask")
plt.axis("off")

/tmp/fix-108932/ipykernel_2324214/2288155435.py:18: UserWarning: Argument(s) 'image_fill_value' are not valid for transform MaskDropout
  A.MaskDropout((1, 1), image_fill_value=255, p=1),

(np.float64(-0.5), np.float64(511.5), np.float64(511.5), np.float64(-0.5))

# The configuration below requires 11GB of VRAM
data_config = {"trainpath": "./data/train", 
               "valid_ratio": 0.2,
               "batch_size": 32,
               "num_workers": 7,
               "patch_size": (256, 256),
               "patch_stride": (128, 128),
               "task": "living_nonliving",
               "normalize": True,
               "max_num_samples": 50000
               }

train_loader, valid_loader, input_size, num_classes, normalizing_stats = data.get_dataloaders(data_config, use_cuda)

# The normalizing statistics must be saved because they will be needed for inference
with open(logdir / "normalizing_stats.yaml", "w") as file:
    yaml.dump(normalizing_stats, file)

100%|██████████| 8/8 [00:00<00:00, 4497.91it/s]
100%|██████████| 1250/1250 [01:24<00:00, 14.83it/s]

print(f"The train dataloader contains {len(train_loader)} mini batches")
print(f"The valid dataloader contains {len(valid_loader)} mini batches")
print(f"Our problem contains {num_classes} classes and our inputs have the shape {input_size}")
print(f"For normalizing the input, we will be using the following statistics {normalizing_stats}")

The train dataloader contains 1250 mini batches
The valid dataloader contains 313 mini batches
Our problem contains 2 classes and our inputs have the shape (1, 256, 256)
For normalizing the input, we will be using the following statistics {'mean': 0.8375427017211914, 'std': 0.09082679715796328}

from planktoseg.models.unet import TimmEncoder

encoder = TimmEncoder(cin=1, model_name="resnet18")
fake_input = torch.zeros(1, 1, 512, 512) # This must be in the pytorch format (B, C, H, W)
f4, [f1, f2, f3] = encoder(fake_input)

print(f"With an input of shape {fake_input.shape}, \n our model outputs tensors of shape :")
print(f" - f1 : {f1.shape}")
print(f" - f2 : {f2.shape}")
print(f" - f3 : {f3.shape}")
print(f" - f4 : {f4.shape}")

With an input of shape torch.Size([1, 1, 512, 512]), 
 our model outputs tensors of shape :
 - f1 : torch.Size([1, 64, 128, 128])
 - f2 : torch.Size([1, 128, 64, 64])
 - f3 : torch.Size([1, 256, 32, 32])
 - f4 : torch.Size([1, 512, 16, 16])

from planktoseg.models.unet import DecoderBlock, Decoder

# First test : we check the forward pass through a decoder block is working
print("First test")
dummy_input = torch.zeros(3, 4, 128, 128)
dummy_encoder_features = torch.zeros(3, 2, 256, 256)
block = DecoderBlock(4)
output = block(dummy_input, dummy_encoder_features)

print(f"With an input of shape {dummy_input.shape} with encoder features of shape {dummy_encoder_features.shape}, the output of the decoder block is of shape {output.shape}")

assert list(output.shape) == [3, 2, 256, 256]

# Second test : we check the forward pass through a complete decoder is working
print("\nSecond test")
batch_size = 3
K = 10
f1 = torch.zeros(batch_size, 64, 128, 128)
f2 = torch.zeros(batch_size, 128, 64, 64)
f3 = torch.zeros(batch_size, 256, 32, 32)
f4 = torch.zeros(batch_size, 512, 16, 16)

decoder = Decoder(num_classes = K)
output = decoder(f4, [f1, f2, f3])

print(f"The output of the decoder is of shape {output.shape}")

assert list(output.shape) == [batch_size, K, 512, 512]
print("If we reach that point, we are good to go !")

First test
With an input of shape torch.Size([3, 4, 128, 128]) with encoder features of shape torch.Size([3, 2, 256, 256]), the output of the decoder block is of shape torch.Size([3, 2, 256, 256])

Second test
The output of the decoder is of shape torch.Size([3, 10, 512, 512])
If we reach that point, we are good to go !

%%capture

model = models.UNet({"encoder": {"model_name": "resnet18"}}, input_size, 1 if num_classes == 2 else num_classes)
model = model.to(device)

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F

gamma = 2
Nsteps = 50
p = torch.linspace(0., 1., steps=Nsteps)
labels = torch.ones_like(p)

bce_loss_values = -torch.log(p)
focal_loss_values = -(1-p)**gamma * torch.log(p)

plt.figure()
plt.plot(p, bce_loss_values, label='BCE loss')
plt.plot(p, focal_loss_values, label='Focal loss')
plt.xlabel("Probability assigned to the class to be predicted")
plt.ylabel("Loss value")
plt.legend()

<matplotlib.legend.Legend at 0x7f6c880ac6d0>

loss = optim.BinaryFocalLoss()

# Optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# Metrics evaluated over the training and validation folds
train_fmetrics = {
    "focal": main.deepcs.metrics.GenericBatchMetric(loss),
    "accuracy": main.BatchAccuracy(),
    "confusion_matrix": metrics.BinaryConfusionMatrixMetric(),
}
train_fmetrics["F1"] = metrics.BinaryF1Metric()
test_fmetrics = {
    "focal": main.deepcs.metrics.GenericBatchMetric(loss),
    "accuracy": main.BatchAccuracy(),
    "confusion_matrix": metrics.BinaryConfusionMatrixMetric(),
}
test_fmetrics["F1"] = metrics.BinaryF1Metric()

# Define the early stopping callback
model_checkpoint = utils.ModelCheckpoint(
    model, logdir, (1,) + input_size, device, min_is_best=False
)

num_epochs = 20

metrics_store = {"train": [], "valid": []}

def postprocess_metrics(fmetrics, metrics_dict):
    """
    This function is used only for extracting
    the precision and recall we track for displaying and ploting
    """
    cm = fmetrics["confusion_matrix"]
    metrics_dict["precision"] = cm.get_precision()
    metrics_dict["recall"] = cm.get_recall()

# Evaluate the metrics before training
train_metrics = utils.test(model, train_loader, device, test_fmetrics)
postprocess_metrics(train_fmetrics, train_metrics)
metrics_store["train"].append(train_metrics)

valid_metrics = utils.test(model, valid_loader, device, test_fmetrics)
postprocess_metrics(test_fmetrics, valid_metrics)
metrics_store["valid"].append(valid_metrics)

for e in range(num_epochs):
        # Train 1 epoch
        train_metrics = utils.train(
            model, train_loader, loss, optimizer, device, train_fmetrics, use_autocast=True, gradient_accumulation=2
        )
        postprocess_metrics(train_fmetrics, train_metrics)
        metrics_store["train"].append(train_metrics)

        # Test
        valid_metrics = utils.test(model, valid_loader, device, test_fmetrics)
        postprocess_metrics(test_fmetrics, valid_metrics)
        metrics_store["valid"].append(valid_metrics)

        # Save the model if it is better
        checkpoint_metric_name = "F1"
        checkpoint_metric = valid_metrics[checkpoint_metric_name]
        updated = model_checkpoint.update(checkpoint_metric)

        # Display the metrics
        metrics_msg = f" Epoch {e} / {num_epochs}\n"
        metrics_msg += "- Train : \n  "
        metrics_msg += "\n  ".join(
            f" {m_name}: {m_value}" for (m_name, m_value) in train_metrics.items()
        )
        metrics_msg += "\n"
        metrics_msg += "- Valid : \n  "
        metrics_msg += "\n  ".join(
            f" {m_name}: {m_value}"
            + ("[>> BETTER <<]" if updated and m_name == checkpoint_metric_name else "")
            for (m_name, m_value) in valid_metrics.items()
        )
        print(metrics_msg)

100%|██████████| 1250/1250 [01:17<00:00, 16.06it/s]
100%|██████████| 313/313 [00:19<00:00, 16.07it/s]
1250it [02:22,  8.80it/s]
100%|██████████| 313/313 [00:19<00:00, 16.00it/s]

 Epoch 0 / 20
- Train : 
   focal: 0.08842981127202511
   accuracy: 0.9569657318115234
   confusion_matrix: [[0.984560818434984, 0.015439181565016006], [0.14412051680573137, 0.8558794831942687]]
   F1: 0.7784466352306805
   precision: 0.7137056591530477
   recall: 0.8558794831942687
- Valid : 
   focal: 0.05701432123780251
   accuracy: 0.9581524917602539
   confusion_matrix: [[0.9918627047558664, 0.008137295244133596], [0.14919320684204732, 0.8508067931579527]]
   F1: 0.8353033018339561[>> BETTER <<]
   precision: 0.8203547115599382
   recall: 0.8508067931579527

1250it [02:21,  8.86it/s]
100%|██████████| 313/313 [00:19<00:00, 16.04it/s]

 Epoch 1 / 20
- Train : 
   focal: 0.03922448422163725
   accuracy: 0.9569657318115234
   confusion_matrix: [[0.9904625031552303, 0.009537496844769644], [0.13008181098011845, 0.8699181890198815]]
   F1: 0.8356533342599359
   precision: 0.8039862340299617
   recall: 0.8699181890198815
- Valid : 
   focal: 0.029857855677604676
   accuracy: 0.9581524917602539
   confusion_matrix: [[0.9932920011982709, 0.006707998801729129], [0.13468788886477076, 0.8653121111352292]]
   F1: 0.8572113073535104[>> BETTER <<]
   precision: 0.8492607716083391
   recall: 0.8653121111352292

1250it [02:20,  8.87it/s]
100%|██████████| 313/313 [00:19<00:00, 16.21it/s]

 Epoch 2 / 20
- Train : 
   focal: 0.024323564728349446
   accuracy: 0.9569657318115234
   confusion_matrix: [[0.9917350603795003, 0.008264939620499722], [0.1359301555745717, 0.8640698444254283]]
   F1: 0.8438503026707409
   precision: 0.8246044818847238
   recall: 0.8640698444254283
- Valid : 
   focal: 0.021647291553020476
   accuracy: 0.9581524917602539
   confusion_matrix: [[0.9908758379932292, 0.009124162006770839], [0.13249373759876096, 0.8675062624012391]]
   F1: 0.8355805587025873
   precision: 0.8059212428402837
   recall: 0.8675062624012391

1250it [02:20,  8.88it/s]
100%|██████████| 313/313 [00:19<00:00, 16.35it/s]

 Epoch 3 / 20
- Train : 
   focal: 0.017641140395402908
   accuracy: 0.9569657318115234
   confusion_matrix: [[0.9920694084442917, 0.007930591555708258], [0.1376783865567481, 0.8623216134432519]]
   F1: 0.8459301773457703
   precision: 0.8302119059302435
   recall: 0.8623216134432519
- Valid : 
   focal: 0.013542211799323559
   accuracy: 0.9581524917602539
   confusion_matrix: [[0.9941612251769757, 0.005838774823024345], [0.1360545160263835, 0.8639454839736165]]
   F1: 0.8649697379272963[>> BETTER <<]
   precision: 0.8659963897237278
   recall: 0.8639454839736165

1250it [02:20,  8.89it/s]
100%|██████████| 313/313 [00:19<00:00, 16.16it/s]

 Epoch 4 / 20
- Train : 
   focal: 0.013638624830171467
   accuracy: 0.9569657318115234
   confusion_matrix: [[0.9926297385773518, 0.007370261422648223], [0.12705171000269547, 0.8729482899973046]]
   F1: 0.857121663716136
   precision: 0.8419291206236177
   recall: 0.8729482899973046
- Valid : 
   focal: 0.011095934394001961
   accuracy: 0.9581524917602539
   confusion_matrix: [[0.9951373790442328, 0.004862620955767134], [0.16323752516072546, 0.8367624748392746]]
   F1: 0.8590556422492047
   precision: 0.8825692001849259
   recall: 0.8367624748392746

import matplotlib.pyplot as plt

plt.figure(figsize=(15,5), dpi=150)
plt.subplot(1, 4, 1)
plt.plot([mi["focal"] for mi in metrics_store["train"]], 'k-')
plt.plot([mi["focal"] for mi in metrics_store["valid"]], 'k--')
plt.title("Focal loss")
plt.xlabel("Epoch")

plt.subplot(1, 4, 2)
plt.plot([mi["F1"] for mi in metrics_store["train"]], 'k-')
plt.plot([mi["F1"] for mi in metrics_store["valid"]], 'k--')
plt.title("F1 score")
plt.xlabel("Epoch")

plt.subplot(1, 4, 3)
plt.plot([mi["precision"] for mi in metrics_store["train"]], 'k-')
plt.plot([mi["precision"] for mi in metrics_store["valid"]], 'k--')
plt.title("Precision")
plt.xlabel("Epoch")

plt.subplot(1, 4, 4)
plt.plot([mi["recall"] for mi in metrics_store["train"]], 'k-')
plt.plot([mi["recall"] for mi in metrics_store["valid"]], 'k--')
plt.title("Recall")
plt.xlabel("Epoch")

Text(0.5, 0, 'Epoch')

import pathlib
import matplotlib.pyplot as plt
import onnxruntime as ort
from PIL import Image
Image.MAX_IMAGE_PIXELS = 25000 * 15000
import numpy as np
import yaml
import glob

providers = []
use_cuda = True
patch_size = 2048

if use_cuda:
    providers.append("CUDAExecutionProvider")
providers.append("CPUExecutionProvider")

# You may adapt the following to either use 
# the model you trained  or the pretrained model you are provided
logdir = pathlib.Path("pretrained_model")
#logdir = pathlib.Path("logs")

inference_session = ort.InferenceSession(
    str(logdir / "best_model.onnx"), providers=providers
)

# Load our normalizing statistics
stats = yaml.safe_load(open(str(logdir / "normalizing_stats.yaml"), "r"))
mean = stats["mean"]
std = stats["std"]

# Load our image
scan_files = list(glob.glob("./data/test/*_scan.png.ppm"))
mask_files = [f.replace("_scan.png.ppm", "_mask.png.ppm") for f in scan_files]
test_idx = 0

scan_path = scan_files[test_idx]
scan_img = np.array(Image.open(scan_path))

mask_path = mask_files[test_idx]
mask_img = np.array(Image.open(mask_path))
print(mask_img.shape)

crop_offset = (5048, 2048)
# Normalize our input
scan_img = ((scan_img - mean * 255.)/(std * 255.)).astype(np.float32)
scan_img = scan_img[np.newaxis, np.newaxis, ...]
scan_img = scan_img[:, :, crop_offset[0]:(crop_offset[0] + patch_size), crop_offset[1]:(crop_offset[1] + patch_size)]

# Get the ground truth mask
# print(np.unique(mask_img))
mask_img = mask_img[crop_offset[0]:(crop_offset[0] + patch_size), crop_offset[1]:(crop_offset[1] + patch_size)] >= 8

# Perform an inference
logits = inference_session.run(None, {"scan": scan_img})[0]
probs = 1.0 / (1.0 + np.exp(-logits))

pred_mask = probs >= 0.5

# Plot the results
plt.figure(dpi=300)
plt.subplot(1, 4, 1)
plt.imshow(scan_img.squeeze(), cmap="gray")
plt.title("Zooscan image")
plt.axis("off")

plt.subplot(1, 4, 2)
plt.imshow(probs.squeeze(), interpolation="none", clim=(0.0, 1.0))
plt.title("Probabilities")
plt.axis("off")

plt.subplot(1, 4, 3)
plt.imshow(pred_mask.squeeze(), interpolation="none", cmap="tab20c")
plt.title("Predicted Mask")
plt.axis("off")

plt.subplot(1, 4, 4)
plt.imshow(mask_img, interpolation="none", cmap="tab20c")
plt.title("Ground truth")
plt.axis("off")

plt.tight_layout()

(14519, 22777)

%pip install segmentation_models_pytorch

Collecting segmentation_models_pytorch
  Downloading segmentation_models_pytorch-0.3.4-py3-none-any.whl (109 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 109.5/109.5 KB 1.4 MB/s eta 0:00:00a 0:00:01
Collecting efficientnet-pytorch==0.7.1
  Downloading efficientnet_pytorch-0.7.1.tar.gz (21 kB)
  Preparing metadata (setup.py) ... done
Requirement already satisfied: huggingface-hub>=0.24.6 in ./venv/lib/python3.10/site-packages (from segmentation_models_pytorch) (0.25.1)
Collecting pretrainedmodels==0.7.4
  Downloading pretrainedmodels-0.7.4.tar.gz (58 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 58.8/58.8 KB 1.6 MB/s eta 0:00:00
  Preparing metadata (setup.py) ... done
Requirement already satisfied: six in ./venv/lib/python3.10/site-packages (from segmentation_models_pytorch) (1.16.0)
Collecting timm==0.9.7
  Downloading timm-0.9.7-py3-none-any.whl (2.2 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.2/2.2 MB 5.1 MB/s eta 0:00:00:00:0100:01
Requirement already satisfied: torchvision>=0.5.0 in ./venv/lib/python3.10/site-packages (from segmentation_models_pytorch) (0.19.1)
Requirement already satisfied: pillow in ./venv/lib/python3.10/site-packages (from segmentation_models_pytorch) (10.4.0)
Requirement already satisfied: tqdm in ./venv/lib/python3.10/site-packages (from segmentation_models_pytorch) (4.66.5)
Requirement already satisfied: torch in ./venv/lib/python3.10/site-packages (from efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (2.4.1)
Collecting munch
  Downloading munch-4.0.0-py2.py3-none-any.whl (9.9 kB)
Requirement already satisfied: safetensors in ./venv/lib/python3.10/site-packages (from timm==0.9.7->segmentation_models_pytorch) (0.4.5)
Requirement already satisfied: pyyaml in ./venv/lib/python3.10/site-packages (from timm==0.9.7->segmentation_models_pytorch) (6.0.2)
Requirement already satisfied: filelock in ./venv/lib/python3.10/site-packages (from huggingface-hub>=0.24.6->segmentation_models_pytorch) (3.16.1)
Requirement already satisfied: typing-extensions>=3.7.4.3 in ./venv/lib/python3.10/site-packages (from huggingface-hub>=0.24.6->segmentation_models_pytorch) (4.12.2)
Requirement already satisfied: packaging>=20.9 in ./venv/lib/python3.10/site-packages (from huggingface-hub>=0.24.6->segmentation_models_pytorch) (24.1)
Requirement already satisfied: fsspec>=2023.5.0 in ./venv/lib/python3.10/site-packages (from huggingface-hub>=0.24.6->segmentation_models_pytorch) (2024.9.0)
Requirement already satisfied: requests in ./venv/lib/python3.10/site-packages (from huggingface-hub>=0.24.6->segmentation_models_pytorch) (2.32.3)
Requirement already satisfied: numpy in ./venv/lib/python3.10/site-packages (from torchvision>=0.5.0->segmentation_models_pytorch) (2.1.1)
Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (12.1.0.106)
Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (12.1.105)
Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (10.3.2.106)
Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (12.1.105)
Requirement already satisfied: jinja2 in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (3.1.4)
Requirement already satisfied: triton==3.0.0 in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (3.0.0)
Requirement already satisfied: sympy in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (1.13.3)
Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (12.1.3.1)
Requirement already satisfied: nvidia-nccl-cu12==2.20.5 in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (2.20.5)
Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (11.4.5.107)
Requirement already satisfied: networkx in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (3.3)
Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (12.1.105)
Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (9.1.0.70)
Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (12.1.105)
Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in ./venv/lib/python3.10/site-packages (from torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (11.0.2.54)
Requirement already satisfied: nvidia-nvjitlink-cu12 in ./venv/lib/python3.10/site-packages (from nvidia-cusolver-cu12==11.4.5.107->torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (12.6.77)
Requirement already satisfied: urllib3<3,>=1.21.1 in ./venv/lib/python3.10/site-packages (from requests->huggingface-hub>=0.24.6->segmentation_models_pytorch) (2.2.3)
Requirement already satisfied: certifi>=2017.4.17 in ./venv/lib/python3.10/site-packages (from requests->huggingface-hub>=0.24.6->segmentation_models_pytorch) (2024.8.30)
Requirement already satisfied: charset-normalizer<4,>=2 in ./venv/lib/python3.10/site-packages (from requests->huggingface-hub>=0.24.6->segmentation_models_pytorch) (3.3.2)
Requirement already satisfied: idna<4,>=2.5 in ./venv/lib/python3.10/site-packages (from requests->huggingface-hub>=0.24.6->segmentation_models_pytorch) (3.10)
Requirement already satisfied: MarkupSafe>=2.0 in ./venv/lib/python3.10/site-packages (from jinja2->torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (2.1.5)
Requirement already satisfied: mpmath<1.4,>=1.1.0 in ./venv/lib/python3.10/site-packages (from sympy->torch->efficientnet-pytorch==0.7.1->segmentation_models_pytorch) (1.3.0)
Using legacy 'setup.py install' for efficientnet-pytorch, since package 'wheel' is not installed.
Using legacy 'setup.py install' for pretrainedmodels, since package 'wheel' is not installed.
Installing collected packages: munch, efficientnet-pytorch, timm, pretrainedmodels, segmentation_models_pytorch
  Running setup.py install for efficientnet-pytorch ... done
  Attempting uninstall: timm
    Found existing installation: timm 1.0.9
    Uninstalling timm-1.0.9:
      Successfully uninstalled timm-1.0.9
  Running setup.py install for pretrainedmodels ... done
Successfully installed efficientnet-pytorch-0.7.1 munch-4.0.0 pretrainedmodels-0.7.4 segmentation_models_pytorch-0.3.4 timm-0.9.7
Note: you may need to restart the kernel to use updated packages.

import segmentation_models_pytorch as smp
import torch

model = smp.DeepLabV3Plus(
    encoder_name="resnet18", 
    encoder_weights="imagenet", 
    in_channels=1, 
    classes=1, 
)

dummy_input = torch.zeros(3, 1, 512, 512)
output = model(dummy_input)
print(f"The output has a shape {output.shape}")

Downloading: "https://download.pytorch.org/models/resnet18-5c106cde.pth" to /usr/users/dce-admin/fix/.cache/torch/hub/checkpoints/resnet18-5c106cde.pth
100%|██████████| 44.7M/44.7M [00:04<00:00, 11.7MB/s]

The output has a shape torch.Size([3, 1, 512, 512])

import wandb
wandb.init(project=..., entity=....)

Classe name	Count
Non living	6.693.443.154
Living	259.647.119

ZooScan image segmentation using a UNet with a pretrained backbone¶

Introduction¶

Imports¶

Data loading and exploration¶

Dataset¶

Exploring the augmentation transforms¶

Adding the transforms in the data pipeline¶

Encoder/Decoder with a pretrained backbone¶

Experimenting with the encoder¶

Implementing the decoder¶

Building the complete model¶

Loss function and metrics for an unbalanced classification¶

Last elements : optimizer, early stopping, metrics, loggers, ...¶

Our first experiments¶

Inference on new zooscan images with non overlapping tiles¶

Going further¶

Using DeepLab v3+ instead of the UNet with pretrained backbone¶

Interfacing with an online dashboard : wandb.ai¶

Running the code without jupyter¶

Classe name	Count
background	6648643689
Chaetognatha	58412864
Salpida	56586684
detritus	33649190
Calanidae	17664122
Aulacantha	12856260
Creseis acicula	8225786
Eumalacostraca	7416376
nectophore<Diphyidae	7036773
badfocus<artefact	6540988
multiple<other	6236479
Rhopalonema velatum	5682772
Calanoida	5602157
Metridinidae	5429665
damaged<Aulacantha	4346996
Chelophyes appendiculata	4061841
nectophore<Physonectae	3743692
Candaciidae	3646939
ephyra<Scyphozoa	3104710
Euchaetidae	2741439
gonophore<Diphyidae	2502463
Abylopsis tetragona	2443476
Doliolida	2442549
Hydrozoa	2418177
Centropagidae	1803662
nectophore<Abylopsis tetragona	1760329
artefact	1672201
othertocheck	1434265
zoea<Galatheidae	1407337
Eucalanidae	1393273
zoea<Brachyura	1360657
Temoridae	1258350
t001	1142346
bract<Abylopsis tetragona	1135173
fiber<detritus	1094546
Aglaura	1071800
Ctenophora<Metazoa	1046977
part<Crustacea	999993
egg<other	995258
bract<Diphyidae	991480
Cavolinia inflexa	954769
tail<Chaetognatha	951886
Creseidae	914248
gonophore<Abylopsis tetragona	908079
Corycaeidae	853293
Oikopleuridae	846252
protozoea<Mysida	780997
calyptopsis<Euphausiacea	780051
bubble	659786
Hyperiidea	570687
eudoxie<Abylopsis tetragona	548613
Sapphirinidae	490605
Heterorhabdidae	487610
Fritillariidae	481907
nucleus<Salpidae	468802
part<Siphonophorae	461031
Liriope<Geryoniidae	428388
Copilia	413538
egg<Actinopterygii	400252
Pyrosomatida	393117
eudoxie<Diphyidae	391068
Actinopterygii	386618
part<Cnidaria	380051
seaweed	361969
Lensia subtilis	331479
larvae<Squillidae	329856
division	326004
head<Chaetognatha	299446
siphonula	286749
part<Mollusca	285145
larvae<Porcellanidae	278677
megalopa	259496
Penilia avirostris	259458
tail<Appendicularia	234629
Solmundella bitentaculata	232568
like<Collodaria	224458
Rhizaria	217177
juvenile<Salpida	201881
pluteus<Ophiuroidea	178319
Annelida	166785
like<Laomediidae	164274
nauplii<Crustacea	133174
Gammaridea	131295
Collodaria	118535
Cymbulia peroni	117545
protozoea<Sergestidae	115471
Gymnosomata	114724
chain<Salpida	109933
Oithonidae	96581
Ostracoda	96258
trunk<Appendicularia	91666
Actiniaria	86592
Haloptilus	84510
egg<Mollusca	76979
Limacinidae	70271
Rhincalanidae	70044
Pontellidae	65340
multiple<Copepoda	64726
Physonectae	61635
colony<Phaeodaria	60429
part<Ctenophora	55696
gonophore<Bassia bassensis	51563
bract<Abylidae	50050
Cliidae	49886
Atlanta	48107
Acartiidae	47749
Amphipoda	45436
pluteus<Echinoidea	45399
t003	40408
phyllosoma	36785
protozoea<Penaeidae	33334
Podon	29154
Obelia	28436
Globigerinidae	24077
Pontellina plumata	22435
Oncaeidae	20112
Orbulina	18316
Insecta	16989
nectophore<Bassia bassensis	13441
nauplii<Cirripedia	9843
Echinodermata	9496
cypris	6894
Evadne	6758
Harpacticoida	5712
t002	4580
Calocalanus pavo	4258
wing	4236
Bivalvia<Mollusca	1663
eudoxie<Bassia bassensis	0
cirrus	0
larvae<Annelida	0
part<Annelida	0
scale	0
nectophore<Hippopodiidae	0
Mollusca	0