def dijkstra(STN, node, string=False, sink=False):
distances = np.zeros(STN.get_num_tp())+np.inf
visited = []
if sink:
nextNodes = STN.get_preds()
else:
nextNodes = STN.get_succs()
if string:
rnode = STN.find_tp(node)
else:
rnode = node
distances[rnode] = 0
p_queue = PriorityQueue()
visited.append(rnode)
p_queue.put((distances[rnode], rnode))
while not p_queue.empty():
node = p_queue.get()[1]
visited.append(node)
for succ in nextNodes[node].keys():
if not succ in visited:
distances[succ] = min(distances[succ], distances[node]+nextNodes[node][succ])
p_queue.put((distances[succ], succ))
if nextNodes[node][succ] < 0:
print('\nDijkstra\'s algorithm can not be used on networks with negative weight edges.\n')
return
return distances
def __init__(self, modelpath=None):
epoch = 30
batch_size = 32
device = 'cpu'
lr = 0.001
kwargs = {'num_workers': 8, 'pin_memory': True}
self.transform = transforms.Compose([
#transforms.Pad(4),
transforms.RandomResizedCrop(28),
#transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
#transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])
#self.KNNModel=self.getKNNmodel(modelpath)
self.SSS = STN()
self.SSS = torch.load(
'/home/wxyice/Desktop/srtp/srtp_code/color_data/colorclassification/PPP.pkl'
)
for name, p in self.SSS.named_parameters():
#print(name)
if name.startswith('localization_linear'): p.requires_grad = False
if name.startswith('localization_convs'): p.requires_grad = False
if name.startswith('convs'): p.requires_grad = False
train_data = DataSTN(modelpath, self.transform)
train_loader = DataLoader(dataset=train_data,
batch_size=batch_size,
shuffle=True,
**kwargs)
train(self.SSS, epoch, lr, train_loader, batch_size, device)
self.SSS.eval()
def floyd_warshall(STN):
num_tp = STN.get_num_tp()
distanceMatrix = np.zeros(shape=(num_tp, num_tp)) + np.inf
successors = STN.get_succs()
for succ_dict in np.arange(num_tp):
distanceMatrix[succ_dict][succ_dict] = 0
for key in successors[succ_dict].keys():
distanceMatrix[succ_dict][key] = successors[succ_dict][key]
for i in np.arange(num_tp):
for j in np.arange(num_tp):
for k in np.arange(num_tp):
distanceMatrix[j][k] = min(distanceMatrix[j][k],
distanceMatrix[j][i] + distanceMatrix[i][k])
return distanceMatrix
def __init__(self, _id, pos_x, pos_y, capability, speed, logger):
self.logger = logger
self.id = _id
self.name = "robot" + str(self.id)
self.init_pos = (pos_x, pos_y)
self.speed = speed
self.stn = STN(self.init_pos, self.speed)
self._bit_schedule = BitSchedule(self.init_pos, self.speed,
self.logger)
self._bid_alpha = 0.5
self._cost_alpha = 0.5
self._capability = capability
self._completed_auctions = []
self._t_auc = []
self._best_task_pos = {}
self._winner_received = True
self._tasks_preconditions = {}
self._auc_ack_pub = None
self._bid_pub = None
self._scheduled_tasks_pub = None
def __init__(self, _id, pos_x, pos_y):
self.init_pos = (pos_x, pos_y)
self.id = int(_id)
self.stn = STN(self.init_pos, 1)
self._bid_alpha = 0.2
self._auc_ack_pub = rospy.Publisher("/auction_ack",
AuctionAck,
queue_size=10)
self._bid_pub = rospy.Publisher("/bid", Bid, queue_size=10)
self._scheduled_tasks_pub = rospy.Publisher("/scheduled_tasks",
ScheduledTasks,
queue_size=10)
self._auction_sub = rospy.Subscriber("/auction", AuctionRequest,
self.auction_callback)
self._winner_sub = rospy.Subscriber("/winner", Winner,
self.winner_callback)
self._completed_auctions = []
self._t_auc = []
self._best_task_pos = {}
self._winner_received = True
self._tasks_preconditions = {}
def __init__(self, input_dim, hidden_dim, kernel_size, in_shape, bias):
"""
Initialize ConvLSTM cell.
Parameters
----------
input_size: (int, int)
Height and width of input tensor as (height, width).
input_dim: int
Number of channels of input tensor.
hidden_dim: int
Number of channels of hidden state.
kernel_size: (int, int)
Size of the convolutional kernel.
bias: bool
Whether or not to add the bias.
"""
super(ConvLSTMCell, self).__init__()
# self.height, self.width = input_size
self.input_dim = input_dim
self.hidden_dim = hidden_dim
self.kernel_size = kernel_size
self.padding = kernel_size[0] // 2, kernel_size[1] // 2
self.bias = bias
self.stn = STN(in_channels=self.input_dim + self.hidden_dim,
in_shape=in_shape)
self.conv = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim,
out_channels=4 * self.hidden_dim,
kernel_size=self.kernel_size,
padding=self.padding,
bias=self.bias)
def test_batch_transform():
transformer = STN(
output_img_size=(32, 100),
num_control_points=6,
margins=[0.1,0.1]
)
test_input_ctrl_pts = np.array([
[
[-0.8, -0.2], [0.0, -0.8], [0.8, -0.2],
[-0.8, 0.8], [0.0, 0.2], [0.8, 0.8]
],
[
[-0.8, -0.8], [0.0, -0.2], [0.8, -0.8],
[-0.8, 0.3], [0.0, 0.8], [0.8, 0.3]
],
[
[-0.8, -0.8], [0.0, -0.8], [0.8, -0.8],
[-0.8, 0.8], [0.0, 0.8], [0.8, 0.8],
]
], dtype=np.float32)
test_im = Image.open('/workspace/xqq/aster/data/test_image.jpg').resize((128, 128))
test_image_array = np.array(test_im).transpose((2,0,1))
test_image_array = np.array([test_image_array, test_image_array, test_image_array])
test_images = torch.Tensor(test_image_array)
test_images = (test_images / 128.0) - 1.0
rectified_images,source_coordinate = transformer(test_images,torch.Tensor(test_input_ctrl_pts))
rectified_images = np.transpose(rectified_images.detach().numpy(),(0,2,3,1))
outputs = {'sampling_grid':source_coordinate.detach().numpy(),'rectified_images':rectified_images}
output_ctrl_pts = transformer.target_control_points
rectified_images_ = (outputs['rectified_images'] + 1.0) * 128.0
if True:
plt.figure()
plt.subplot(3, 4, 1)
plt.scatter(test_input_ctrl_pts[0, :, 0], test_input_ctrl_pts[0, :, 1])
plt.subplot(3, 4, 2)
plt.scatter(output_ctrl_pts[:, 0], output_ctrl_pts[:, 1])
plt.subplot(3, 4, 3)
plt.scatter(outputs['sampling_grid'][0, :, 0], outputs['sampling_grid'][0, :, 1], marker='+')
plt.subplot(3, 4, 4)
plt.imshow(rectified_images_[0].astype(np.uint8))
plt.subplot(3, 4, 5)
plt.scatter(test_input_ctrl_pts[1, :, 0], test_input_ctrl_pts[1, :, 1])
plt.subplot(3, 4, 6)
plt.scatter(output_ctrl_pts[:, 0], output_ctrl_pts[:, 1])
plt.subplot(3, 4, 7)
plt.scatter(outputs['sampling_grid'][1, :, 0], outputs['sampling_grid'][1, :, 1], marker='+')
plt.subplot(3, 4, 8)
plt.imshow(rectified_images_[1].astype(np.uint8))
plt.subplot(3, 4, 9)
plt.scatter(test_input_ctrl_pts[2, :, 0], test_input_ctrl_pts[2, :, 1])
plt.subplot(3, 4, 10)
plt.scatter(output_ctrl_pts[:, 0], output_ctrl_pts[:, 1])
plt.subplot(3, 4, 11)
plt.scatter(outputs['sampling_grid'][2, :, 0], outputs['sampling_grid'][2, :, 1], marker='+')
plt.subplot(3, 4, 12)
plt.imshow(rectified_images_[2].astype(np.uint8))
plt.show()
class Robot:
def __init__(self, _id, pos_x, pos_y):
self.init_pos = (pos_x, pos_y)
self.id = int(_id)
self.stn = STN(self.init_pos, 1)
self._bid_alpha = 0.2
self._auc_ack_pub = rospy.Publisher("/auction_ack",
AuctionAck,
queue_size=10)
self._bid_pub = rospy.Publisher("/bid", Bid, queue_size=10)
self._scheduled_tasks_pub = rospy.Publisher("/scheduled_tasks",
ScheduledTasks,
queue_size=10)
self._auction_sub = rospy.Subscriber("/auction", AuctionRequest,
self.auction_callback)
self._winner_sub = rospy.Subscriber("/winner", Winner,
self.winner_callback)
self._completed_auctions = []
self._t_auc = []
self._best_task_pos = {}
self._winner_received = True
self._tasks_preconditions = {}
def __str__(self):
s = "id: " + str(self.id)
s += ", init location: " + str(self.init_pos)
return s
def start_listener(self):
rospy.init_node("robot" + str(self.id))
rospy.spin()
def auction_callback(self, msg):
if msg.id in self._completed_auctions:
return
auc_ack_msg = AuctionAck()
auc_ack_msg.robot_id = int(self.id)
auc_ack_msg.auc_id = int(msg.id)
#print "Robot " + str(self.id) + ": Sending ack for auction " + str(msg.id)
print "Robot " + str(self.id) + ": Received auction " + str(
msg.id) + " with following " + str(len(msg.tasks)) + " tasks."
self._auc_ack_pub.publish(auc_ack_msg)
self._t_auc = utils.create_tasks(msg.tasks)
for t in self._t_auc:
print str(t)
i = 0
for task in self._t_auc:
self._tasks_preconditions[task] = msg.PC[i]
i += 1
self._winner_received = True
while len(self._t_auc) > 0:
min_bid = float("inf")
min_task = None
if self._winner_received:
print "Robot " + str(self.id) + ": Finding optimal task"
for task in self._t_auc:
bid, best_pos = self._compute_min_bid(task)
self._best_task_pos[task.id] = best_pos
if bid < min_bid:
min_bid = bid
min_task = task
if min_task is not None:
print "Robot " + str(self.id) + ": Found optimal task"
bid_msg = Bid()
bid_msg.auc_id = msg.id
bid_msg.robot_id = self.id
bid_msg.task = utils.create_task_msg(min_task)
bid_msg.bid = min_bid
print "Robot " + str(self.id) + ": Publishing bid"
self._bid_pub.publish(bid_msg)
self._winner_received = False
tasks = self._tighten_schedule()
print "Robot " + str(
self.id) + "All tasks from the current auction has been scheduled."
print "Robot " + str(self.id) + ": Makespan is " + str(
self.stn.get_makespan())
print "Robot " + str(self.id) + ": Schedule: " + str(self.stn)
scheduled_tasks_msg = ScheduledTasks()
scheduled_tasks_msg.robot_id = self.id
scheduled_tasks_msg.tasks = tasks
self._scheduled_tasks_pub.publish(scheduled_tasks_msg)
self._completed_auctions.append(msg.id)
def winner_callback(self, msg):
print "Robot " + str(self.id) + ": Winner received"
winner_robot_id = msg.robot_id
task = utils.create_task(msg.task)
if winner_robot_id == self.id:
print "Robot " + str(self.id) + " is the winner"
self.stn.insert_task(task, self._best_task_pos[task.id])
self.stn.solve_stn(self._tasks_preconditions)
print "Robot " + str(self.id) + ": Following task has been scheduled."
print str(task)
print ""
self._t_auc.remove(task)
self._winner_received = True
def _tighten_schedule(self):
print "Robot " + str(self.id) + ": Tightening schedule"
tasks = []
for i in range(self.stn.task_count):
task = self.stn._get_task(i)
f = task.finish_time
task.update_time_window(task.est, f)
self.stn.update_task_constraints(task.id)
tasks.append(task)
return tasks
def _compute_min_bid(self, task):
task_count = self.stn.task_count
min_bid = float("inf")
min_pos = None
for i in range(task_count + 1):
temp_stn = deepcopy(self.stn)
tt_before = temp_stn.total_travel_time
temp_stn.insert_task(task, i)
temp_stn.solve_stn(self._tasks_preconditions)
tt_after = temp_stn.total_travel_time
if temp_stn.is_consistent():
addition_travel_time = tt_after - tt_before
bid = self._compute_bid(temp_stn.get_makespan(),
addition_travel_time)
if bid < min_bid:
min_bid = bid
min_pos = i
return min_bid, min_pos
def _compute_bid(self, makespan, travel_time):
return makespan * self._bid_alpha + travel_time * (1 - self._bid_alpha)
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Oct 9 08:37:40 2017
@author: jeff
"""
import sys
sys.path.append('../STN')
from STN import STN
# create instance
stn = STN()
# states
stn.state('FeedA', init=200)
stn.state('FeedB', init=200)
stn.state('FeedC', init=200)
stn.state('HotA', price=-1)
stn.state('IntAB', price=-1)
stn.state('IntBC', price=-1)
stn.state('ImpureE', price=-1)
stn.state('Product_1', price=10)
stn.state('Product_2', price=10)
# state to task arcs
stn.stArc('FeedA', 'Heating')
stn.stArc('FeedB', 'Reaction_1', rho=0.5)
stn.stArc('FeedC', 'Reaction_1', rho=0.5)
from cluttered_mnist import ClutteredMNIST
from STN import STN
from paz.backend.image import write_image
dataset_path = "mnist_cluttered_60x60_6distortions.npz"
batch_size = 256
num_epochs = 10
save_path = ''
data_manager = ClutteredMNIST(dataset_path)
train_data, val_data, test_data = data_manager.load()
x_train, y_train = train_data
model = STN()
model.compile(loss={'label': 'categorical_crossentropy'}, optimizer='adam')
model.summary()
def plot_predictions(samples):
(lables, interpolations) = model.predict(samples)
for arg, images in enumerate(zip(interpolations, samples)):
interpolated, image = images
interpolated = (interpolated * 255).astype('uint8')
image = (image * 255).astype('uint8')
write_image('images/interpolated_image_%03d.png' % arg, interpolated)
write_image('images/original_image_%03d.png' % arg, image)
model.fit(x_train, y_train, batch_size, num_epochs, validation_data=val_data)
plot_predictions(test_data[0][:9])
class Robot():
def __init__(self, _id, pos_x, pos_y, capability, speed, logger):
self.logger = logger
self.id = _id
self.name = "robot" + str(self.id)
self.init_pos = (pos_x, pos_y)
self.speed = speed
self.stn = STN(self.init_pos, self.speed)
self._bit_schedule = BitSchedule(self.init_pos, self.speed,
self.logger)
self._bid_alpha = 0.5
self._cost_alpha = 0.5
self._capability = capability
self._completed_auctions = []
self._t_auc = []
self._best_task_pos = {}
self._winner_received = True
self._tasks_preconditions = {}
self._auc_ack_pub = None
self._bid_pub = None
self._scheduled_tasks_pub = None
def __str__(self):
s = "id: " + str(self.id)
s += ", init location: " + str(self.init_pos)
return s
def set_alpha(self, alpha):
self._bid_alpha = alpha
self._cost_alpha = alpha
def start_listener(self):
rospy.init_node(self.name, disable_signals=True)
self._auc_ack_pub = rospy.Publisher("/auction_ack",
AuctionAck,
queue_size=10)
self._bid_pub = rospy.Publisher("/bid", Bid, queue_size=10)
self._scheduled_tasks_pub = rospy.Publisher("/scheduled_tasks",
ScheduledTasks,
queue_size=10)
rospy.Subscriber("/auction", AuctionRequest, self.auction_callback)
rospy.Subscriber("/winner", Winner, self.winner_callback)
rospy.Service(self.name + '/terminate', TerminateRobot,
self.handle_robot_terminate)
rospy.spin()
return self.stn.get_makespan(), self.stn.total_travel_time
def auction_callback(self, msg):
if msg.id in self._completed_auctions:
return
auc_ack_msg = AuctionAck()
auc_ack_msg.robot_id = int(self.id)
auc_ack_msg.auc_id = int(msg.id)
self._auc_ack_pub.publish(auc_ack_msg)
self.logger.debug(
"Robot {0}: Received auction {1} with following {2} tasks.".format(
self.id, msg.id, len(msg.tasks)))
self._t_auc = utils.create_tasks(msg.tasks)
for t in self._t_auc:
self.logger.debug("{0}".format(str(t)))
i = 0
for task in self._t_auc:
self._tasks_preconditions[task] = msg.PC[i]
i += 1
self._winner_received = True
while len(self._t_auc) > 0:
min_bid = float("inf")
min_task = None
if self._winner_received:
self.logger.info("Robot {0}: Finding optimal task".format(
self.id))
for task in self._t_auc:
bid, best_pos = self._compute_min_bid(task)
self.logger.debug(
"Robot {0}: Bid for task {1} is {2}, and best pos is {3}"
.format(self.id, task.id, bid, best_pos))
self._best_task_pos[task.id] = best_pos
if bid < min_bid:
min_bid = bid
min_task = task
if min_task is not None:
self.logger.info("Robot {0}: Found optimal task".format(
self.id))
bid_msg = Bid()
bid_msg.auc_id = msg.id
bid_msg.robot_id = self.id
bid_msg.task = utils.create_task_msg(min_task)
bid_msg.bid = min_bid
self.logger.info("Robot {0}: Publishing bid {1}".format(
self.id, str(bid_msg.task)))
self._bid_pub.publish(bid_msg)
self._winner_received = False
tasks = self.tighten_schedule()
self.logger.info(
"Robot {0}: All tasks from the current auction has been scheduled."
.format(self.id))
self.logger.debug("Robot {0}: Makespan is {1}".format(
self.id, self.stn.get_makespan()))
self.logger.debug("\nRobot {0}: Schedule:\n {1}\n".format(
self.id, str(self.stn)))
self._publish_scheduled_tasks(tasks)
self._completed_auctions.append(msg.id)
def winner_callback(self, msg):
self.logger.info("Robot {0}: Winner received".format(self.id))
winner_robot_id = msg.robot_id
task = utils.create_task(msg.task)
if winner_robot_id == self.id:
self.logger.info("Robot {0} is the winner".format(self.id))
self.stn.insert_task(task, self._best_task_pos[task.id])
self.stn.solve_stn(self._tasks_preconditions)
self.logger.debug(
"Robot {0}: Following task has been scheduled.".format(self.id))
self.logger.debug("{0}".format(str(task)))
self._t_auc.remove(task)
self._winner_received = True
def handle_robot_terminate(self, req):
rospy.signal_shutdown("Normal Shutdown")
return TerminateRobotResponse(True)
def init_auction(self, auc_id, t_auct, PC):
self._t_auc = deepcopy(t_auct)
self.logger.debug("Robot {0}: Initializing Auction {1}".format(
self.id, auc_id))
for t in self._t_auc:
self.logger.debug("{0}".format(str(t)))
i = 0
for task in self._t_auc:
self._tasks_preconditions[task] = PC[i]
i += 1
def get_min_bid(self):
min_bid = float("inf")
min_task = None
self.logger.info("Robot {0}: Finding optimal task".format(self.id))
for task in self._t_auc:
bid, best_pos = self._compute_min_bid(task)
self.logger.debug(
"Robot {0}: Bid for task {1} is {2}, and best pos is {3}".
format(self.id, task.id, bid, best_pos))
self._best_task_pos[task.id] = best_pos
if bid < min_bid:
min_bid = bid
min_task = task
return min_bid, min_task
def notify_winner(self, auc_id, winner_robot_id, task):
self.logger.info("Robot {0}: Winner received".format(self.id))
if winner_robot_id == self.id:
self.logger.info("Robot {0} is the winner".format(self.id))
self.stn.insert_task(task, self._best_task_pos[task.id])
self.stn.solve_stn(self._tasks_preconditions)
self.logger.debug(
"Robot {0}: Following task has been scheduled.".format(self.id))
self.logger.debug("{0}".format(str(task)))
self._t_auc.remove(task)
def end_auction(self):
self._t_auc = []
self.logger.warn("Robot {0}: Ending auction".format(self.id))
def is_capable(self, task):
if task.type not in self._capability:
return False
return True
def get_best_cost(self, task, pc): # pc -precondition or previous tasks
task_count = self.stn.task_count
min_cost = float("inf")
min_pos = None
for i in range(task_count + 1):
task_copy = deepcopy(task)
stn_copy = deepcopy(self.stn)
tt_before = stn_copy.total_travel_time
stn_copy.insert_task(task_copy, i)
stn_copy.solve_stn(pc)
tt_after = stn_copy.total_travel_time
if stn_copy.is_consistent():
addition_travel_time = tt_after - tt_before
ms = stn_copy.get_makespan()
cost = self._compute_cost(ms, addition_travel_time)
if cost < min_cost:
min_cost = cost
min_pos = i
return min_cost, min_pos
def get_ms_tt(self, task, pc, insert_time):
ms = None
addition_travel_time = None
task_copy = deepcopy(task)
stn_copy = deepcopy(self.stn)
tt_before = stn_copy.total_travel_time
stn_copy.insert_task(task_copy, time=insert_time)
stn_copy.solve_stn(pc)
tt_after = stn_copy.total_travel_time
if stn_copy.is_consistent():
addition_travel_time = tt_after - tt_before
ms = stn_copy.get_makespan()
return (ms, addition_travel_time)
def get_cost(self, task, pc, insert_time):
cost = None
ms, tt = get_ms_tt(task, pc, insert_time)
if ms != None and tt != None:
cost = self._compute_cost(ms, tt)
return cost
def add_task(self, task, pc, pos=None, time=None):
if (pos is None) and (time is None):
self.logger(
"Robot {0} ADD TASK: Either pos or time needs to be given, not both."
.format(self.id))
return
# print("during adding1")
# print(self.stn)
#Here the new task is not added yet
if time is not None:
# print("robot", self.id, "start inserting using time:", time)
self.stn.insert_task(task, time=time)
else:
# print("robot", self.id, "start inserting using pos:", pos)
self.stn.insert_task(task, index=pos)
# print("during adding2")
# print(self.stn)
#Here the new task is added but start and end times are 0
self.stn.solve_stn(pc)
# print("during adding3")
# print(self.stn)
#Here start and end time is updated
self._bit_schedule = BitSchedule(self.init_pos,
self.speed,
self.logger,
stn=self.stn)
# print("during adding4")
# print(self.stn)
def get_bit_schedule(self, new_task):
schedule_copy = deepcopy(self._bit_schedule)
if new_task is None:
return schedule_copy
schedule_copy.prepare_for_coalition(new_task)
return schedule_copy
def _publish_scheduled_tasks(self, tasks):
scheduled_tasks_msg = ScheduledTasks()
scheduled_tasks_msg.robot_id = self.id
scheduled_tasks_msg.tasks = tasks
scheduled_tasks_msg.makespan = self.stn.get_makespan()
scheduled_tasks_msg.travel_time = self.stn.total_travel_time
self._scheduled_tasks_pub.publish(scheduled_tasks_msg)
def tighten_schedule(self):
self.logger.info("Robot {0}: Tightening schedule".format(self.id))
tasks = []
for i in range(self.stn.task_count):
task = self.stn._get_task(i)
f = task.finish_time
task.update_time_window(task.est, f)
self.stn.update_task_constraints(task.id)
tasks.append(task)
return tasks
def _compute_min_bid(self, task):
task_count = self.stn.task_count
min_bid = float("inf")
min_pos = None
for i in range(task_count + 1):
stn_copy = deepcopy(self.stn)
tt_before = stn_copy.total_travel_time
stn_copy.insert_task(task, i)
stn_copy.solve_stn(self._tasks_preconditions)
tt_after = stn_copy.total_travel_time
if stn_copy.is_consistent():
addition_travel_time = tt_after - tt_before
bid = self._compute_bid(stn_copy.get_makespan(),
addition_travel_time)
if bid < min_bid:
min_bid = bid
min_pos = i
return min_bid, min_pos
def _compute_bid(self, makespan, travel_time):
return makespan * self._bid_alpha + travel_time * (1 - self._bid_alpha)
def _compute_cost(self, makespan, travel_time):
return makespan * self._cost_alpha + travel_time * (1 -
self._cost_alpha)
def prop_fwd_prop_bkwd(STN, edge, string=True):
if string:
x = STN.find_tp(edge[0])
delta = int(edge[1])
y = STN.find_tp(edge[2])
else:
x = edge[0]
delta = edge[1]
y = edge[2]
if not STN.get_dist_mat_upd():
print('You must have an updated distance matrix')
quit()
# get distance matrix, succ, preds
dist_mat = STN.get_dist_mat()
successors = STN.get_succs()
predecessors = STN.get_preds()
# if existing distance less than new edge weight, return current dist matrix
if dist_mat[x][y] <= delta:
return dist_mat
# update distances
successors[x][y] = delta
predecessors[y][x] = delta
dist_mat[x][y] = delta
# keep track during fwd prop
encountered = [y]
changed = [y]
to_do = [y]
if delta < -dist_mat[y][x]:
return False
while to_do:
v = to_do[0]
to_do.remove(v)
v_succs = successors[v]
for w in v_succs.keys():
d = v_succs[w]
if (not w in encountered) and (dist_mat[y][w] == d + dist_mat[y][v]):
encountered.append(w)
if delta+dist_mat[y][w] < dist_mat[x][w]:
dist_mat[x][w] = delta+dist_mat[y][w]
changed.append(w)
to_do.append(w)
for w in changed:
encountered = [x]
to_do = [x]
while to_do:
f = to_do[0]
to_do.remove(f)
f_preds = predecessors[f]
for e in f_preds.keys():
a = f_preds[e]
if (not e in encountered) and (dist_mat[e][x] == a + dist_mat[f][x]):
encountered.append(e)
if dist_mat[e][x] + dist_mat[x][w] < dist_mat[e][w]:
dist_mat[e][w] = dist_mat[e][x] + dist_mat[x][w]
to_do.append(e)
return dist_mat
class KNNClassifier:
def __init__(self, modelpath=None):
epoch = 30
batch_size = 32
device = 'cpu'
lr = 0.001
kwargs = {'num_workers': 8, 'pin_memory': True}
self.transform = transforms.Compose([
#transforms.Pad(4),
transforms.RandomResizedCrop(28),
#transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
#transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])
#self.KNNModel=self.getKNNmodel(modelpath)
self.SSS = STN()
self.SSS = torch.load(
'/home/wxyice/Desktop/srtp/srtp_code/color_data/colorclassification/PPP.pkl'
)
for name, p in self.SSS.named_parameters():
#print(name)
if name.startswith('localization_linear'): p.requires_grad = False
if name.startswith('localization_convs'): p.requires_grad = False
if name.startswith('convs'): p.requires_grad = False
train_data = DataSTN(modelpath, self.transform)
train_loader = DataLoader(dataset=train_data,
batch_size=batch_size,
shuffle=True,
**kwargs)
train(self.SSS, epoch, lr, train_loader, batch_size, device)
self.SSS.eval()
def getKNNmodel(self, modelpath):
KNNModel = []
for i in os.listdir(modelpath):
if i != ".DS_Store":
for j in os.listdir(os.path.join(modelpath, i)):
if j != ".DS_Store":
img = cv.imread(os.path.join(modelpath, i, j))
m = main_color_moment(img)
k = [m, str(i)]
KNNModel.append(k)
return KNNModel
def prediction(self, img, k=5):
if self.transform:
img = Image.fromarray(img)
img = self.transform(img)
img = img.view(1, 3, 28, 28)
output = self.SSS(img)
pred = output.data.max(1, keepdim=True)[1]
return pred
# m=main_color_moment(img)
# distancelist=[]
# pset=set()
# for i in self.KNNModel:
# d=self.distance(m,i[0])
# label=i[-1]
# distancelist.append([d,label])
# pset.add(label)
# distancelist=sorted(distancelist)
# pdict={}
# for i in pset:
# pdict[i]=0
# for i in range(k):
# pdict[distancelist[i][1]]+=1
# return classes_name.index(max(pdict, key=pdict.get))
def distance(self, p1, p2):
# vector1=np.array(p1)
# vector2=np.array(p2)
#op1=np.sqrt(np.sum(np.square(vector1-vector2)))
try:
op2 = np.linalg.norm(p1 - p2)
except Exception as Error:
print('[distance wrong]' + Error)
return op2