-
Notifications
You must be signed in to change notification settings - Fork 1
/
graspPredMod.py
114 lines (99 loc) · 5.02 KB
/
graspPredMod.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
76
77
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
#!/usr/bin/env python
import numpy as np
import cv2
import copy
from sys import stdout
# Given image, returns image point and theta to grasp
class Predictors:
def __init__(self, I, learner=None):
self.I = I
self.I_h, self.I_w, self.I_c = self.I.shape
self.learner = learner
def sigmoid_array(self, x):
return 1 / (1 + np.exp(-x))
def to_prob_map(self, fc8_vals):
sig_scale = 1
no_keep = 20
fc8_sig = self.sigmoid_array(sig_scale * fc8_vals)
r = np.sort(fc8_sig, axis=None)
r_no_keep = r[-no_keep]
fc8_sig[fc8_sig < r_no_keep] = 0.0
fc8_prob_map = fc8_sig / fc8_sig.sum()
return fc8_prob_map
def sample_from_map(self, prob_map):
prob_map_contig = np.ravel(prob_map)
arg_map_contig = np.array(range(prob_map_contig.size))
smp = np.random.choice(arg_map_contig, p=prob_map_contig)
return np.unravel_index(smp, prob_map.shape)
def random_grasp(self, num_angle=18):
h_g = np.random.randint(self.I_h)
w_g = np.random.randint(self.I_w)
t_g = np.random.randint(num_angle)
return h_g, w_g, t_g
def center_grasp(self, num_angle=18):
h_g = np.int(self.I_h / 2)
w_g = np.int(self.I_w / 2)
t_g = np.int(num_angle / 2)
return h_g, w_g, t_g
def graspNet_grasp(self, init_loc, num_angle=18, patch_size=300, num_samples=128):
self.patch_size = patch_size
half_patch_size = np.int(patch_size / 2) + 1
h_range = self.I_h - patch_size - 2
w_range = self.I_w - patch_size - 2
# Initialize random patch points
patch_hs = np.random.randint(
h_range, size=num_samples) + half_patch_size
patch_ws = np.random.randint(
w_range, size=num_samples) + half_patch_size
# print 'initialized patch_hs is: ', patch_hs
# print 'initialized patch_ws is: ', patch_ws
patch_Is = np.zeros((num_samples, patch_size, patch_size, self.I_c))
print "patch_Is initialized with shape", patch_Is.shape
# print "patch_Is initialized as", patch_Is
patch_Is_resized = np.zeros(
(num_samples, self.learner.IMAGE_SIZE, self.learner.IMAGE_SIZE, self.I_c))
for looper in xrange(num_samples):
isWhiteFlag = 1
while isWhiteFlag == 1:
# patch_hs[looper] = np.random.randint(h_range) + half_patch_size
# patch_ws[looper] = np.random.randint(w_range) + half_patch_size
patch_hs[looper] = init_loc[1] + np.random.randint(low=-10, high=10, size=1)
patch_ws[looper] = init_loc[0] + np.random.randint(low=-10, high=10, size=1)
patch_Is[looper] = self.I[patch_hs[looper] - half_patch_size:patch_hs[looper] - half_patch_size +
patch_size,
patch_ws[looper] - half_patch_size:patch_ws[looper] - half_patch_size + patch_size]
# Make sure that the input has a minimal amount of standard deviation from mean. If
# not resample
if patch_Is[looper].std() > 15:
isWhiteFlag = 0
else:
isWhiteFlag = 1
patch_Is_resized[looper] = cv2.resize(
patch_Is[looper], (self.learner.IMAGE_SIZE, self.learner.IMAGE_SIZE), interpolation=cv2.INTER_CUBIC)
# subtract mean
patch_Is_resized = patch_Is_resized - 111
self.fc8_vals = self.learner.test_one_batch(patch_Is_resized)
# Normalizing angle uncertainity
wf = [0.25, 0.5, 0.25]
self.fc8_norm_vals = copy.deepcopy(self.fc8_vals)
for looper in xrange(num_samples):
for norm_looper in xrange(num_angle):
self.fc8_norm_vals[looper, norm_looper] = (wf[1] * self.fc8_norm_vals[looper, norm_looper] +
wf[0] * self.fc8_norm_vals[
looper, (norm_looper - 1) % num_angle] +
wf[2] * self.fc8_norm_vals[
looper, (norm_looper + 1) % num_angle])
# Normalize to probability distribution
self.fc8_prob_vals = self.to_prob_map(self.fc8_norm_vals)
# Sample from probability distribution
self.patch_id, self.theta_id = self.sample_from_map(self.fc8_prob_vals)
# self.patch_id, self.theta_id = np.unravel_index(self.fc8_prob_vals.argmax(), self.fc8_prob_vals.shape)
self.patch_hs = patch_hs
self.patch_ws = patch_ws
self.best_patch = patch_Is[self.patch_id]
self.best_patch_h = patch_hs[self.patch_id]
self.best_patch_w = patch_ws[self.patch_id]
self.best_patch_resized = patch_Is_resized[self.patch_id]
self.patch_Is_resized = patch_Is_resized
self.patch_Is = patch_Is
return self.best_patch_h, self.best_patch_w, self.theta_id