seperate into app.py and segmentation.py

app.py

@@ -8,68 +8,17 @@ import cv2
 from PIL import Image as PIL_Image
 from io import BytesIO
 import matplotlib.pyplot as plt
+from segmentation import ImageSegmentation
 
-# def flip_text(x):
-#     return x[::-1]
-# def flip_image(x):
-#     return np.fliplr(x)
 
 def image_segmentation(x, threshold_factor):
-    cv2.startWindowThread()
-    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()
-
-    # img = cv2.imread(x) # 改為使用Gradio傳入的圖片
-    img = x
-    # plt.imshow(img[..., ::-1])
-
-    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
-    ret, 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)
-    ret, sure_fg = cv2.threshold(dist_transform,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
-    ret, 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(img,markers)
-    img[markers == -1] = [255,0,0]
-
-    # plt.imshow(dist_transform)
-    # plt.imshow(sure_fg)
-    # plt.imshow(sure_bg, alpha=0.3)
-    # plt.imshow(markers)
-
-    # 將 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  # 返回彩色分割結果圖
+    segmetation = ImageSegmentation(x, threshold_factor)
+    output_img = segmetation.segmented_img
+    return output_img
+    
 
 with gr.Blocks() as demo:
     gr.Markdown("Image Segmentation using Gradio")
-    # with gr.Tab("Flip Text"):
-    #     text_input = gr.Textbox()
-    #     text_output = gr.Textbox()
-    #     text_button = gr.Button("Flip")
     with gr.Tab("Image Segmentation"):
         with gr.Row():
             image_input = gr.Image()
@@ -83,17 +32,6 @@ with gr.Blocks() as demo:
         )
         image_button = gr.Button("Flip")
 
-    # with gr.Accordion("Open for More!", open=False):
-    #     gr.Markdown("Look at me...")
-    #     temp_slider = gr.Slider(
-    #         0, 1,
-    #         value=0.1,
-    #         step=0.1,
-    #         interactive=True,
-    #         label="Slide me",
-    #     )
-
-    # text_button.click(flip_text, inputs=text_input, outputs=text_output)
     image_button.click(image_segmentation, inputs=[image_input, temp_slider], outputs=image_output)
 
 if __name__ == "__main__":

segmentation.py

@@ -7,41 +7,49 @@ from PIL import Image as PIL_Image
 from io import BytesIO
 import matplotlib.pyplot as plt
 
-cv2.startWindowThread()
-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()
+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)
 
-img = cv2.imread("img/water_coins.jpg") # 改為使用Gradio傳入的圖片
-# plt.imshow(img[..., ::-1])
+        # 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]
 
-gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
-ret, 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)
+        # 將 markers 範圍縮放到 0-255 並轉換為 uint8 類型
+        markers_scaled = (markers - markers.min()) / (markers.max() - markers.min()) * 255
+        markers_scaled = markers_scaled.astype(np.uint8)
 
-# sure background area
-sure_bg = cv2.dilate(opening,kernel,iterations=3)
-# Finding sure foreground area
-dist_transform = cv2.distanceTransform(opening,cv2.DIST_L2,5)
-ret, sure_fg = cv2.threshold(dist_transform,0.7*dist_transform.max(),255,0)
-# Finding unknown region
-sure_fg = np.uint8(sure_fg)
-unknown = cv2.subtract(sure_bg,sure_fg)
-# Marker labelling
-ret, 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(img,markers)
-img[markers == -1] = [255,0,0]
+        # 使用 applyColorMap 生成彩色圖像
+        markers_colored = cv2.applyColorMap(markers_scaled, cv2.COLORMAP_JET)
 
-# plt.imshow(dist_transform)
-# plt.imshow(sure_fg)
-# plt.imshow(sure_bg, alpha=0.3)
-# plt.imshow(markers)
+        return markers_colored  # 返回彩色分割結果圖