Incorrect label assignment after watershed segmentation (Python)

I am segmenting some particles using skimage.segmentation.watershed. This works quite well, and I manage to separate them adecuately (see image). Image after segmentation. However, when I use ndi.label to label the different regions, not all the segments get separated, and some are assigned the same label (even if with watershed they were correctly segmented) (image 2, for instance top left particle) after labeling assignment. Notice how particles previously separated are assigned the same label. I am relatively new in Python and do not know what else I can try to solve this. I would appreciate if you could give me any help. Thanks in advance :)

The code I am using is (I have it in a for loop since I want to automatize the process for the analysis of several images simultaneously):

#import hyperspy for reading directly ser or emd files
import hyperspy.api as hs
#The scikit image library will be used for segmenting the images
from skimage.exposure import histogram
from skimage.color import label2rgb
from skimage import data, io, filters
from skimage.filters import threshold_local, threshold _yen, threshold_li
from skimage.filters import try_all_threshold
from skimage.filters import gaussian
from skimage.feature import peak_local_max
from skimage.feature import canny
from skimage import measure
from skimage.morphology import label
from skimage.morphology import remove_small_objects
from skimage.draw import ellipse
from skimage.measure import label, regionprops, regionprops_table
from skimage.transform import rotate
from skimage.segmentation import watershed
#matplotlib for performing plots
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

#Basic packages (math and statistics)
import pandas as pd
import numpy as np
from scipy import ndimage as ndi
from scipy import stats
import math
import glob
import seaborn as sns

#load data
s=hs.load(folder+'*.emi',stack=True)

#threshold
thresh=threshold_li(s.data)
binary=s>thresh

#Cleaning

cleaned=remove_small_objects(binary.data, min_size=5)

Segmentation itself

#Define variables needed

dist=np.zeros([cleaned.shape[1],cleaned.shape[1]])
water=np.zeros([cleaned.shape[0],cleaned.shape[1],cleaned.shape[1]])
mask=np.zeros(dist.shape, dtype=bool)
markers,_=ndi.label(mask)
water_particles=np.zeros([cleaned.shape[0],cleaned.shape[1],cleaned.shape[1]])
eq_diam_total=np.array([])

#for loop for segmenting all the images using watershed. 
#I will use the commented "for i in range (cleaned.shape[0])" 
#once I manage to solve the segmentation issue:


#for i in range (cleaned.shape[0]):
for i in range(2,3):
    dist = ndi.distance_transform_edt(cleaned[i,:,:]) #make distance map
    maxima=peak_local_max(gaussian(dist, sigma=1.5),threshold_rel=None, 
    min_distance=5)  # find local maxima
    print('maxima',maxima.shape)

    mask[tuple(maxima.T)]=True
    markers,_=ndi.label(mask)
    print('markers',markers.shape)
 
    #segment
    water[i,:,:]=watershed(-dist, markers,mask=cleaned[i,:,:])
    print('water', water.shape)
    
    #label each particle
    water_particles[i,:,:], water_labels = ndi.label(water[i,:,:])
    print('water_particles',water_particles.shape)
    print('water_labels',water_labels)
  
    

Plot after segmentation

%matplotlib inline  
from skimage import color
fig,axes=plt.subplots(1, 2, sharey=True)  

axes[0].imshow(color.label2rgb(water[i,:,:]))
axes[0].axis('off')
axes[0].set_title('After watershed segmentation')               
               

axes[1].imshow(color.label2rgb(water_particles[i,:,:]))
axes[1].axis('off')
axes[1].set_title('After label')

The output of skimage.segmentation.watershed is a labeled image, it can go directly into the measurement function (regionprops).

Also, that watershed function will take all local minima as markers, you don’t need to find these yourself. Simplified code:

from scipy import ndimage as ndi
from skimage.segmentation import watershed

# ...

    dist = ndi.distance_transform_edt(cleaned[i,:,:]) #make distance map
    water_particles[i,:,:] = watershed(-dist, markers,mask=cleaned[i,:,:])

If you want to receive an image you can label yourself (for whatever reason), then add the watershed_line=True argument. This will keep a separation between each region (basin) so that the labeling algorithm can identify them.