Segmenting annotated image

Atif_seg.py

import numpy as np
from PIL import Image
import time

start = time.time()

# Find bounding box using numpy
# https://stackoverflow.com/questions/31400769/bounding-box-of-numpy-array
def bbox2(img):
    rows = np.any(img, axis=1)
    cols = np.any(img, axis=0)
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]

    return rmin, rmax, cmin, cmax

maskim = Image.open("P03 -instance_mask.tiff")
rawim = Image.open("P03.tiff")
maskarr = np.array(maskim,dtype=np.uint16)
rawarr = np.array(rawim,dtype=np.uint8)

unique_colours = np.unique(maskarr)
# Discard black background
unique_colours = [c for c in unique_colours if c!=0]

for colour in unique_colours:
    # Find pixels belonging to specific fibre
    fibmask = maskarr==colour
    bbox = bbox2(fibmask)

    # Crop arrays to bounding box
    fibarr = np.copy(rawarr[bbox[0]:bbox[1],bbox[2]:bbox[3],:])
    
    if(fibarr.size>0):
        fibmask2 = fibmask[bbox[0]:bbox[1],bbox[2]:bbox[3]]

        # Delete any non-fibre pixels that remain in bounding box
        fibarr[np.logical_not(fibmask2)] = [0,0,0]

        # Convert cropped array to image and save to disk
        fibim = Image.fromarray(fibarr, mode="RGB")
        fibim = fibim.resize((fibim.size[0]*5,fibim.size[1]*5),Image.NEAREST)
        fibim.save(f'fibre_{colour:04d}.png',quality=75)

end = time.time()
total_time = end - start
print("\n"+ str(total_time))