← Back to branch summary

~ci-train-bot/ubuntu-settings-components/ubuntu-settings-components-ubuntu-zesty-2329

« back to all changes in this revision

Viewing changes to plugins/Ubuntu/Settings/Fingerprint/utils/segment.py

Merging with trunk

Show diffs side-by-side

added added

removed removed

Lines of Context:
plugins/Ubuntu/Settings/Fingerprint/utils/segment.py
 
1
#!/usr/bin/python2
 
2
"""
 
3
 * Copyright 2016 Canonical Ltd.
 
4
 *
 
5
 * This program is free software; you can redistribute it and/or modify
 
6
 * it under the terms of the GNU Lesser General Public License as published by
 
7
 * the Free Software Foundation; version 3.
 
8
 *
 
9
 * This program is distributed in the hope that it will be useful,
 
10
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 
11
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
12
 * GNU Lesser General Public License for more details.
 
13
 *
 
14
 * You should have received a copy of the GNU Lesser General Public License
 
15
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
16
 *
 
17
 * Authored by  Florian Boucault <florian.boucault@canonical.com>
 
18
 */
 
19
"""
 
20
 
 
21
import os.path
 
22
import argparse
 
23
import json
 
24
import cv2
 
25
 
 
26
parser = argparse.ArgumentParser(
 
27
    description='Generates segmented version of img and extracts bbs.'
 
28
)
 
29
parser.add_argument('input_image', type=str,
 
30
                    help='path to the image to segment')
 
31
args = parser.parse_args()
 
32
 
 
33
input_image = os.path.abspath(args.input_image)
 
34
print("Segmenting %s..." % input_image)
 
35
filename, _ = os.path.splitext(os.path.basename(input_image))
 
36
split = filename.split('@')
 
37
if len(split) != 1:
 
38
    output_image = "%s_segmented@%s.png" % (split[0], split[1])
 
39
    output_boxes = "%s_boxes.json" % split[0]
 
40
else:
 
41
    output_image = "%s_segmented.png" % filename
 
42
    output_boxes = "%s_boxes.json" % filename
 
43
 
 
44
img = cv2.imread(input_image, -1)
 
45
img_boxes = img.copy()
 
46
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
 
47
ret, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)
 
48
contours, hierarchy = cv2.findContours(
 
49
    thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
 
50
)
 
51
 
 
52
print("Detected %d segments." % len(contours))
 
53
 
 
54
bounding_boxes = []
 
55
for index, contour in enumerate(contours):
 
56
    x, y, w, h = cv2.boundingRect(contour)
 
57
    bounding_boxes.append((x, y, w, h))
 
58
#  print(x, y, w, h)
 
59
    cv2.rectangle(img_boxes,
 
60
                  (x, y),
 
61
                  (x + w, y + h),
 
62
                  (255, 20, 20, 255), 1)
 
63
    cv2.drawContours(thresh, contours, index, (index + 1),
 
64
                     cv2.cv.CV_FILLED)
 
65
    cv2.drawContours(thresh, contours, index, (index + 1), 2)
 
66
 
 
67
b, g, r, a = cv2.split(img)
 
68
img = cv2.merge([a, g, thresh])
 
69
 
 
70
cv2.imwrite(output_image, img)
 
71
print("Written segmented version of '%s' to '%s'" % (
 
72
    args.input_image, output_image)
 
73
)
 
74
 
 
75
description = {"width": img.shape[1],
 
76
               "height": img.shape[0],
 
77
               "boxes": bounding_boxes}
 
78
with open(output_boxes, 'w') as outfile:
 
79
    json.dump(description, outfile)
 
80
print("Written bounding boxes to '%s'" % output_boxes)
 
81
 
 
82
cv2.imshow('image', img_boxes)
 
83
cv2.waitKey(0)
 
84
cv2.destroyAllWindows()