chromosomes.py

#!/usr/bin/env python3
from PIL import Image, ImageDraw, ImageChops
import numpy as np
import colorsys

def distinct_colours(n):
    return [colorsys.hsv_to_rgb(i / (n - 1), 1, 1) for i in range(n)]

def in_bounds(point, bounds):
    x, y = point
    w, h = bounds
    return (0 <= x < w) and (0 <= y < h) 

def is_dark(colour):
    return colour < 10

def autocrop(im, bgcolour):
    bg = Image.new(im.mode, im.size, bgcolour)
    diff = ImageChops.difference(im, bg)
    bbox = diff.getbbox()
    if bbox:
        return im.crop(bbox)
    return None

img_orig = Image.open("Patient1.jpg")
pixels_orig = img_orig.load()

arr = np.zeros(shape=img_orig.size, dtype=np.int16)

stack = []
idx = 0
for y in range(img_orig.size[1]):
    for x in range(img_orig.size[0]):
        if is_dark(pixels_orig[x, y]) or arr[x, y]:
            continue
        
        idx += 1

        if idx > 50:
            break

        stack.append((x, y))
        while stack:
            i, j = stack.pop()
            if arr[i, j] or is_dark(pixels_orig[i, j]):
                continue
            arr[i, j] = idx

            stack.extend([(i + ii, j + jj) for ii in range(-1, 2) for jj in range(-1, 2) if in_bounds((i + ii, j + jj), img_orig.size)]) 

print(idx)
colours = [(int(r * 255), int(g * 255), int(b * 255)) for r, g, b in distinct_colours(idx + 1)]

chromosomes = []

for val in range(1, idx + 1):
    indices = np.array((arr == val).nonzero())
    min_x, max_x = np.min(indices[0]), np.max(indices[0])
    min_y, max_y = np.min(indices[1]), np.max(indices[1])
    w, h = max_x - min_x + 1, max_y - min_y + 1

    x = indices[0]
    y = indices[1]
    A = np.vstack([x, np.ones(len(x))]).T
    m, c = np.linalg.lstsq(A, y, rcond=None)[0]
    angle = np.degrees(np.arctan(m))

    
    chromosome = Image.new(img_orig.mode, (w, h), "black")
    chromosome_pixels = chromosome.load()
    for i in range(indices.shape[1]):
        x, y = indices[:,i]
        chromosome_pixels[int(x - min_x), int(y - min_y)] = pixels_orig[int(x), int(y)]
    
    chromosome = autocrop(chromosome.rotate(90 + angle, expand=True, resample=Image.BILINEAR), "black")
    if chromosome.size[0] > chromosome.size[1]:
        chromosome = autocrop(chromosome.rotate(90, expand=True), "black")
    chromosomes.append(chromosome)

chromosomes.sort(key=lambda im: -im.size[1])

for val, chromosome in enumerate(chromosomes):
    chromosome.save(f"chr/chr{val + 1:02d}.png")


img = Image.new("RGB", img_orig.size, "black")
pixels = img.load()

for y in range(img.size[1]):
    for x in range(img.size[0]):
        if arr[x, y]:
            pixels[x, y] = colours[arr[x, y] - 1]

img.show()