Slic.py 6.13 KB
Newer Older
1 2 3 4 5 6 7 8 9 10
#!/usr/bin/python
# -*- coding: utf-8 -*-
#
"""
  Nome: Slic.py
  Autor: Alessandro dos Santos Ferreira ( santosferreira.alessandro@gmail.com )

  Descricão: TODO.
"""

11
import cv2
12 13 14
import numpy as np
from skimage.segmentation import slic
from skimage.segmentation import mark_boundaries
15 16
from skimage.util import img_as_float, img_as_ubyte
from collections import OrderedDict
17 18 19

from util.Config import Config
from util.Utils import TimeUtils
20
from util.X11Colors import Colors
21 22 23 24 25

from Segmentation import Segmentation

class Slic(Segmentation):
    
26
    n_segments = None
27 28
    sigma = None
    compactness = None
29 30
    border_color = None
    border_outline = None
31
    
32 33 34 35 36 37
    __segments = None
    
    __original_image = None
    
    def __init__(self, n_segments = 100, sigma = 5.0, compactness = 10.0, border_color = 'Yellow', border_outline = 'No'):
        self.n_segments = Config("Segments", n_segments, int)
38 39
        self.sigma = Config("Sigma", sigma, float)
        self.compactness = Config("Compactness", compactness, float)
40 41
        self.border_color = Config("Border Color", border_color, 'color')
        self.border_outline = Config("Border Outline", border_outline, str)
42 43 44
        
    
    def get_config(self):
45 46 47 48 49 50 51 52 53
        slic_config = OrderedDict()
        
        slic_config["n_segments"] = self.n_segments
        slic_config["sigma"] = self.sigma
        slic_config["compactness"] = self.compactness
        slic_config["border_color"] = self.border_color
        slic_config["border_outline"] = self.border_outline
        
        return slic_config
54 55
        
    def set_config(self, configs):
56
        self.n_segments = Config.nvl_config(configs["n_segments"], self.n_segments)
57 58
        self.sigma = Config.nvl_config(configs["sigma"], self.sigma)
        self.compactness = Config.nvl_config(configs["compactness"], self.compactness)
59 60 61 62
        self.border_color = Config.nvl_config(configs["border_color"], self.border_color)
        self.border_outline = Config.nvl_config(configs["border_outline"], self.border_outline)
        
        self.border_outline.value = self.border_outline.value if self.border_outline.value == 'Yes' else 'No'
63 64

    def get_summary_config(self):
65 66 67 68 69 70 71 72 73 74 75 76 77
        slic_config = OrderedDict()
        
        slic_config[self.n_segments.label] = self.n_segments.value
        slic_config[self.sigma.label] = self.sigma.value
        slic_config[self.compactness.label] = self.compactness.value
        slic_config[self.border_color.label] = self.border_color.value
        slic_config[self.border_outline.label] = self.border_outline.value
        
        summary = ''
        for config in slic_config:
            summary += "%s: %s\n" % (config, str(slic_config[config]))
        
        return summary
78 79
    

80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150
    def get_segment(self, px = 0, py = 0, idx_segment = None):
        if self.__segments is None:
            return None, 0, -1, 0
        
        start_time = TimeUtils.get_time()
        
        if idx_segment is None:
            idx_segment = self.__segments[py, px]
        
        mask_segment = np.zeros(self.__original_image.shape[:2], dtype="uint8")
        mask_segment[self.__segments == idx_segment] = 255
        size_segment = mask_segment[self.__segments == idx_segment].size

        segment = self.__original_image.copy()
        segment = cv2.bitwise_and(segment, segment, mask=mask_segment)

        contours, _  = cv2.findContours(mask_segment,cv2.RETR_LIST,cv2.CHAIN_APPROX_SIMPLE)[-2:]

        m = -1
        max_contour = None
        for cnt in contours:
            if (len(cnt) > m):
                m = len(cnt)
                max_contour = cnt

        x,y,w,h = cv2.boundingRect(max_contour)
        segment = segment[y:y+h, x:x+w]
        
        end_time = TimeUtils.get_time()
        
        return segment, size_segment, idx_segment, (end_time - start_time)
        
        
    def paint_segment(self, image, color, px = 0, py = 0, idx_segment = None, border = True, clear = False):
        
        if self.__segments is None:
            return image, 0
            
        start_time = TimeUtils.get_time()
        
        if idx_segment is None:
            idx_segment = self.__segments[py, px]
        height, width, channels = self.__original_image.shape
        
        mask_segment = np.zeros(self.__original_image.shape[:2], dtype="uint8")
        mask_segment[self.__segments == idx_segment] = 255
        mask_inv = cv2.bitwise_not(mask_segment)
            
        class_color = np.zeros((height,width,3), np.uint8)
        class_color[:, :] = Colors.get_color(color)
        if clear == False:
            colored_image = cv2.addWeighted(self.__original_image, 0.7, class_color, 0.3, 0)
        else:
            colored_image = self.__original_image
        colored_image = cv2.bitwise_and(colored_image, colored_image, mask=mask_segment)
        
        new_image = cv2.bitwise_and(image, image, mask=mask_inv)
        mask_segment[:] = 255
        new_image = cv2.bitwise_or(new_image, colored_image, mask=mask_segment)

        if border == True:
            color = Colors.get_color_zero_one(self.border_color.get_cast_val())
            outline_color = color if self.border_outline.value == 'Yes' else None
            
            new_image = img_as_ubyte( mark_boundaries(img_as_float(new_image), self.__segments.astype(np.int8), color=color, outline_color=outline_color) ) 
            
        end_time = TimeUtils.get_time()
        
        return new_image, (end_time - start_time)
        

151 152
    def run(self, image):
        
153 154 155 156 157
        self.__original_image = image
        
        color = Colors.get_color_zero_one(self.border_color.get_cast_val())
        outline_color = color if self.border_outline.value == 'Yes' else None
        
158
        start_time = TimeUtils.get_time()
159 160
        self.__segments = slic(img_as_float(image), 
                            n_segments=self.n_segments.get_cast_val(), 
161 162 163 164
                            sigma=self.sigma.get_cast_val(), 
                            compactness=self.compactness.get_cast_val())
        end_time = TimeUtils.get_time()
        
165 166
        return img_as_ubyte( mark_boundaries(image, self.__segments.astype(np.int8), color=color, outline_color=outline_color) ), (end_time - start_time)