-
Notifications
You must be signed in to change notification settings - Fork 3
/
preprocessing_utils.py
75 lines (62 loc) · 2.16 KB
/
preprocessing_utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
import sys
try:
print('removed ros path')
sys.path.remove('/opt/ros/kinetic/lib/python2.7/dist-packages')import cv2 as cv
except:
pass
import numpy as np
import argparse
import os
import json
import csv
def disp_mask(image = None, Mask = None):
alpha = 0.5
try:
image = args.image
mask = args.mask
print('--------------------------------')
except:
pass
image = cv.imread(image)
mask = np.load(mask)
if image is None:
print("Error loading src1")
exit(-1)
elif mask is None:
print("Error loading src2")
exit(-1)
mask[:,:,1] = 0;
seg = cv.addWeighted(image, alpha, mask, (1-alpha), 0.0)
cv.imshow('Segmented', seg)
cv.waitKey(0)
cv.destroyAllWindows()
def get_mask():
if not os.path.exists(args.directory):
os.mkdir(args.directory)
print('Created out_dir')
data = csv.reader(open('data.csv'))
for name in os.listdir(args.input_image):
frame = cv.imread(args.input_image + '/' + name)
with open(str('./' + args.input_annotations + '/' + name.strip('.jpg') + '.json')) as json_file:
data = json.load(json_file)
key_points = data['shapes'][0]['points']
mask = np.zeros(frame.shape, np.uint8)
cv.fillPoly(mask, np.int32([np.array(key_points)]), color = (255,255,255))
# cv.imshow('temp',mask)
# cv.waitKey(0)
# cv.destroyAllWindows()
cv.imwrite(args.directory + '/' + name, mask)
if __name__ == "__get_mask__":
parser = argparse.ArgumentParser()
parser.add_argument("-d", "--directory", type=str, help="Path to save the images")
parser.add_argument("-ia", "--input_annotations", type=str, help="Bag file to read")
parser.add_argument("-ii", "--input_image", type=str, help="Bag file to read")
args = parser.parse_args()
main()
if __name__ == "__disp_mask__":
parser = argparse.ArgumentParser()
parser.add_argument("-i", "--image", type=str, help="path/image")
parser.add_argument("-m", "--mask", type=str, help="path/mask")
parser.add_argument("-a", "--alpha", type=int, help="alpha")
args = parser.parse_args()
main()