Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import cv2
- import numpy as np
- import glob
- import math
- import StringIO
- from scipy.signal import argrelmax, find_peaks
- # ============================================================================
- zones_a = [(50, 72), (144, 166), (243, 265), (328, 350), (436, 458)]
- zones_b = [(42, 64), (122, 144), (207, 229), (276, 298), (369, 391), (496, 518)]
- # ============================================================================
- def plot_report(filename, report):
- from matplotlib import pyplot as plt
- from matplotlib.gridspec import GridSpec
- IMAGE_KEYS = ['img', 'thresh', 'thresh_1', 'thresh_2', 'canny', 'canny_1', 'canny_2']
- PLOT_SPAN = 5
- TEXT_SPAN = 2
- ROW_COUNT = (len(IMAGE_KEYS) + 1) + 3 * (PLOT_SPAN + 1) + 2 * (TEXT_SPAN)
- fig = plt.figure()
- plt.suptitle(filename)
- gs = GridSpec(ROW_COUNT, 2)
- row = 0
- for key in IMAGE_KEYS:
- plt.subplot(gs[row,:])
- plt.gca().set_title(key)
- plt.imshow(report[key], cmap='gray', aspect='equal')
- plt.axis('off')
- row += 1
- proj_width = len(report['projection'])
- proj_x = np.arange(proj_width)
- plt.subplot(gs[row+1:row+1+PLOT_SPAN,:])
- plt.gca().set_title('Vertical Projections (Raw and Smoothed)')
- plt.plot(proj_x, report['projection'], 'y-')
- plt.plot(proj_x, report['projection_1'], 'r-')
- plt.plot(proj_x, report['projection_2'], 'b-')
- plt.xlim((0, proj_width - 1))
- plt.ylim((0, 255))
- row += PLOT_SPAN + 1
- #bins_ab, weight_ab, details_ab, checkbox_count
- plt.subplot(gs[row+1:row+1+PLOT_SPAN,:])
- plt.gca().set_title('Smoothed Projection with Peaks and Zones')
- plt.plot(proj_x, report['projection_2'])
- for zone in zones_a:
- plt.axvspan(zone[0], zone[1], facecolor='y', alpha=0.1)
- for zone in zones_b:
- plt.axvspan(zone[0], zone[1], facecolor='r', alpha=0.1)
- for x in report['peaks']:
- plt.axvline(x=x, color='m')
- plt.xlim((0, proj_width - 1))
- plt.ylim((0, report['projection_2'].max()))
- row += PLOT_SPAN + 1
- plt.subplot(gs[row+1:row+1+TEXT_SPAN,0], frameon=False)
- plt.gca().set_title('Details - 5 boxes')
- plt.axis([0, 1, 0, 1])
- plt.gca().axes.get_yaxis().set_visible(False)
- plt.gca().axes.get_xaxis().set_visible(False)
- plt.text(0, 1, report['details_a'], family='monospace', fontsize=8, ha='left', va='top')
- plt.subplot(gs[row+1:row+1+TEXT_SPAN,1], frameon=False)
- plt.gca().set_title('Details - 6 boxes')
- plt.axis([0, 1, 0, 1])
- plt.gca().axes.get_yaxis().set_visible(False)
- plt.gca().axes.get_xaxis().set_visible(False)
- plt.text(0, 1, report['details_b'], family='monospace', fontsize=8, ha='left', va='top')
- row += TEXT_SPAN
- plt.subplot(gs[row+1:row+1+PLOT_SPAN,:])
- plt.gca().set_title('Weights')
- plt.barh([2, 1]
- , [report['weight_a'], report['weight_b']]
- , align='center'
- , color=['y', 'r']
- , tick_label=['5 boxes', '6 boxes'])
- plt.ylim((0.5, 2.5))
- row += PLOT_SPAN + 1
- row += 1
- plt.subplot(gs[row,:])
- plt.gca().set_title('Input Image')
- plt.imshow(report['img'], cmap='gray', aspect='equal')
- plt.axis('off')
- row += 1
- plt.subplot(gs[row:row+TEXT_SPAN,:], frameon=False)
- plt.axis([0, 1, 0, 1])
- plt.gca().axes.get_yaxis().set_visible(False)
- plt.gca().axes.get_xaxis().set_visible(False)
- result_text = "The image contains %d boxes." % report['checkbox_count']
- plt.text(0.5, 1, result_text, family='monospace', weight='semibold', fontsize=24, ha='center', va='top')
- fig.set_size_inches(12, ROW_COUNT * 0.8)
- plt.savefig('plot_%s.png' % filename[:2], bbox_inches="tight")
- plt.close(fig)
- def bin_peaks(peaks, values, zones):
- bins = [[] for x in xrange(len(zones))]
- for peak, value in zip(peaks, values):
- for i, zone in enumerate(zones):
- if (peak >= zone[0]) and (peak <= zone[1]):
- bins[i].append((peak, value))
- return bins
- def analyze_bins(bins):
- buf = StringIO.StringIO()
- total_weight = 0.0
- for i, bin in enumerate(bins):
- buf.write("Position %d: " % i)
- weight = 0.0
- if len(bin) == 0:
- buf.write("no peaks")
- else:
- best_bin = sorted(bin, key=lambda x: x[1], reverse=True)[0]
- weight = best_bin[1]
- if len(bin) == 1:
- buf.write("single peak @ %d (value=%0.3f)" % best_bin)
- else:
- buf.write("%d peaks, best @ %d (value=%0.3f)" % (len(bin), best_bin[0], best_bin[1]))
- buf.write(" | weight=%0.3f\n" % weight)
- total_weight += weight
- buf.write("Total weight = %0.3f" % total_weight)
- return total_weight, buf.getvalue()
- def proc(filename):
- report = {}
- img = cv2.imread(filename, cv2.IMREAD_GRAYSCALE)
- report['img'] = img.copy()
- thresh = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 15, 2)
- report['thresh'] = thresh.copy()
- thresh = cv2.morphologyEx(thresh, cv2.MORPH_ERODE, np.ones((1,3),np.uint8))
- report['thresh_1'] = thresh.copy()
- thresh = cv2.morphologyEx(thresh, cv2.MORPH_DILATE, np.ones((3,1),np.uint8))
- report['thresh_2'] = thresh.copy()
- edges = cv2.Canny(thresh, 40, 120, apertureSize=5)
- report['canny'] = edges.copy()
- edges = cv2.morphologyEx(edges, cv2.MORPH_OPEN, np.ones((5,1),np.uint8))
- report['canny_1'] = edges.copy()
- edges = cv2.morphologyEx(edges, cv2.MORPH_DILATE, np.ones((1,3),np.uint8))
- report['canny_2'] = edges.copy()
- projection = np.mean(edges, 0).flatten()
- report['projection'] = projection.copy()
- projection = cv2.blur(projection, (1, 21)).flatten()
- report['projection_1'] = projection.copy()
- projection = cv2.blur(projection, (1, 11)).flatten()
- report['projection_2'] = projection.copy()
- peaks = find_peaks(projection)[0]
- report['peaks'] = peaks.copy()
- peak_values = projection[peaks]
- report['peak_values'] = peak_values.copy()
- bins_a = bin_peaks(peaks, peak_values, zones_a)
- report['bins_a'] = list(bins_a)
- bins_b = bin_peaks(peaks, peak_values, zones_b)
- report['bins_b'] = list(bins_b)
- weight_a, details_a = analyze_bins(bins_a)
- report['weight_a'] = weight_a
- report['details_a'] = details_a
- weight_b, details_b = analyze_bins(bins_b)
- report['weight_b'] = weight_b
- report['details_b'] = details_b
- checkbox_count = 5 if (weight_a > weight_b) else 6
- report['checkbox_count'] = checkbox_count
- return checkbox_count, report
- for filename in glob.glob('*-*.png'):
- box_count, report = proc(filename)
- plot_report(filename, report)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement