Advertisement
Guest User

Untitled

a guest
Oct 4th, 2017
75
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 3.33 KB | None | 0 0
  1. # -*- coding: utf-8 -*-
  2.  
  3. import os, sys
  4. import numpy as np
  5. from scipy import ndimage as ndi
  6. from scipy.misc import imsave
  7. import matplotlib.pyplot as plt
  8.  
  9. from skimage.filters import sobel, threshold_local
  10. from skimage.morphology import watershed
  11. from skimage import io
  12.  
  13.  
  14. def open_image(name):
  15.     filename = os.path.join(os.getcwd(), name)
  16.     return io.imread(filename, as_grey=True)
  17.  
  18.  
  19. def adaptive_threshold(image):
  20.  
  21.     # Create threshold image
  22.     # Offset is not desirable for these images
  23.     block_size = 701
  24.     threshold_img = threshold_local(image, block_size)
  25.  
  26.     # Binarize the image with the threshold image
  27.     binary_adaptive = image < threshold_img
  28.  
  29.     # Convert the mask (which has dtype bool) to dtype int
  30.     # This is required for the code in `segmentize` (below) to work
  31.     binary_adaptive = binary_adaptive.astype(int)  
  32.  
  33.     # Return the binarized image
  34.     return binary_adaptive
  35.  
  36.  
  37. def segmentize(image):
  38.     # make segmentation using edge-detection and watershed
  39.     edges = sobel(image)
  40.     markers = np.zeros_like(image)
  41.     foreground, background = 1, 2
  42.     markers[image == 0] = background
  43.     markers[image == 1] = foreground
  44.  
  45.     ws = watershed(edges, markers)
  46.  
  47.     return ndi.label(ws == foreground)
  48.  
  49.  
  50. def find_segment(segments, index):
  51.     segment = np.where(segments == index)
  52.     shape = segments.shape
  53.  
  54.     minx, maxx = max(segment[0].min() - 1, 0), min(segment[0].max() + 1, shape[0])  
  55.     miny, maxy = max(segment[1].min() - 1, 0), min(segment[1].max() + 1, shape[1])
  56.  
  57.     im = segments[minx:maxx, miny:maxy] == index
  58.  
  59.     return (np.sum(im), np.invert(im))
  60.  
  61.  
  62. def run(f):
  63.     print 'Processing:', f
  64.  
  65.     image = open_image(f)
  66.     processed = adaptive_threshold(image)
  67.     segments = segmentize(processed)
  68.  
  69.     print 'Segments detected:', segments[1]
  70.    
  71.     seg = []
  72.     for s in range(1, segments[1]):
  73.         seg.append(find_segment(segments[0], s))
  74.  
  75.     seg.sort(key=lambda s: -s[0])
  76.  
  77.     # Get the directory name (if a full path is given)
  78.     folder = r'Users\J\Desktop\sketch2\data'
  79.  
  80.     # Get the file name
  81.     filenm = os.path.basename(f)
  82.  
  83.     # If it doesn't already exist, create a new dir "segments"
  84.     # to save the PNGs
  85.     segments_folder = os.path.join(folder,filenm[:-4]+"_segments")
  86.     os.path.isdir(segments_folder) or os.mkdir(segments_folder)
  87.    
  88.     # Load color palette and extract colors
  89.     palette = io.imread(os.path.join(os.getcwd(), 'palette.png'))
  90.     palette = palette.reshape(palette.shape[0]*palette.shape[1],palette.shape[2])
  91.     palette_colors = np.unique(palette,axis=0)
  92.     p = palette_colors.shape[0]
  93.    
  94.     # Save the segments to the "segments" directory
  95.     for i in range(len(seg)):
  96.        
  97.          # Create an MxNx4 array (RGBA)
  98.         seg_rgba = np.zeros((seg[i][1].shape[0],seg[i][1].shape[1],4), dtype=np.uint8)
  99.  
  100.         # Fill R, G and B with appropriately colored copies of the image
  101.         for c in range(3):
  102.             seg_rgba[:,:,c] = ~seg[i][1] * palette_colors[i % p,c]
  103.  
  104.         # For A (alpha), use the invert of the image (so background is 0=transparent)
  105.         seg_rgba[:,:,3] = (~seg[i][1]).astype(np.uint8) * 255
  106.  
  107.         # Save image
  108.         imsave(os.path.join(segments_folder, filenm[:-4] + '_seg' + str(i) + '_col' + str(i % p) + '.png'), seg_rgba)
  109.  
  110. for f in sys.argv[1:]:
  111.     run(f)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement