import numpy as np
import cv2
from PIL import Image as PIL_Image
from io import BytesIO

class ImageSegmentation():
    def __init__(self, x, threshold_factor):
        self.img = x
        self.threshold_factor = threshold_factor
        self.segmented_img = self.segmentation()
        
    def img_to_png(ima, cvt=None):
        if cvt:
            ima = cv2.cvtColor(ima, cvt)
        im = PIL_Image.fromarray(ima)
        bio = BytesIO()
        im.save(bio, format='png')
        return bio.getvalue()
    
    def segmentation(self):
        cv2.startWindowThread()
        gray = cv2.cvtColor(self.img,cv2.COLOR_BGR2GRAY)
        _, thresh = cv2.threshold(gray,0,255,cv2.THRESH_BINARY_INV+cv2.THRESH_OTSU)
        kernel = np.ones((3,3),np.uint8)
        opening = cv2.morphologyEx(thresh,cv2.MORPH_OPEN,kernel, iterations = 2)

        # sure background area
        sure_bg = cv2.dilate(opening,kernel,iterations=3)
        # Finding sure foreground area
        dist_transform = cv2.distanceTransform(opening,cv2.DIST_L2,5)
        _, sure_fg = cv2.threshold(dist_transform,self.threshold_factor*dist_transform.max(),255,0)
        # Finding unknown region
        sure_fg = np.uint8(sure_fg)
        unknown = cv2.subtract(sure_bg,sure_fg)
        # Marker labelling
        _, markers0 = cv2.connectedComponents(sure_fg)
        # Add one to all labels so that sure background is not 0, but 1
        markers = markers0+1
        # Now, mark the region of unknown with zero
        markers[unknown==255] = 0
        markers = cv2.watershed(self.img,markers)
        self.img[markers == -1] = [255,0,0]

        # 將 markers 範圍縮放到 0-255 並轉換為 uint8 類型
        markers_scaled = (markers - markers.min()) / (markers.max() - markers.min()) * 255
        markers_scaled = markers_scaled.astype(np.uint8)

        # 使用 applyColorMap 生成彩色圖像
        markers_colored = cv2.applyColorMap(markers_scaled, cv2.COLORMAP_JET)

        return markers_colored  # 返回彩色分割結果圖