idx = opt.idx
print("model %d/%d" % (idx, len(d)))
point, seg = d[idx]
print(point.size(), seg.size())
point_np = point.numpy()
cmap = plt.cm.get_cmap("hsv", 10)
cmap = np.array([cmap(i) for i in range(10)])[:, :3]
gt = cmap[seg.numpy() - 1, :]
classifier = PointNetDenseCls(k=4)
classifier.load_state_dict(torch.load(opt.model))
classifier.eval()
point = point.transpose(1, 0).contiguous()
point = Variable(point.view(1, point.size()[0], point.size()[1]))
pred, _ = classifier(point)
pred_choice = pred.data.max(2)[1]
# print(pred_choice)
#print(pred_choice.size())
pred_color = cmap[pred_choice.numpy()[0], :]
# print(point_np.shape)
# print(pred_color.shape)
pred_color = pred_color * 255
pred_color = pred_color.astype(int)
def main(argv=None):
print('Hello! This is XXXXXX Program')
## Load PointNet config
parser = argparse.ArgumentParser()
parser.add_argument('--model', type=str, default='./seg/seg_model_1.pth', help='model path')
opt = parser.parse_args()
print(opt)
## Load PointNet model
num_points = 2700
classifier = PointNetDenseCls(num_points=num_points, k=10)
classifier.load_state_dict(torch.load(opt.model))
classifier.eval()
### Config visualization
cmap = plt.cm.get_cmap("hsv", 5)
cmap = np.array([cmap(i) for i in range(10)])[:, :3]
# gt = cmap[seg - 1, :]
## Initialize OpenNi
# dist = './driver/OpenNI-Linux-x64-2.3/Redist'
dist = './driver/OpenNI-Windows-x64-2.3/Redist'
openni2.initialize(dist)
if (openni2.is_initialized()):
print("openNI2 initialized")
else:
print("openNI2 not initialized")
## Register the device
dev = openni2.Device.open_any()
## Create the streams stream
rgb_stream = dev.create_color_stream()
depth_stream = dev.create_depth_stream()
## Define stream parameters
w = 320
h = 240
fps = 30
## Configure the rgb_stream -- changes automatically based on bus speed
rgb_stream.set_video_mode(
c_api.OniVideoMode(pixelFormat=c_api.OniPixelFormat.ONI_PIXEL_FORMAT_RGB888, resolutionX=w, resolutionY=h,
fps=fps))
## Configure the depth_stream -- changes automatically based on bus speed
# print 'Depth video mode info', depth_stream.get_video_mode() # Checks depth video configuration
depth_stream.set_video_mode(
c_api.OniVideoMode(pixelFormat=c_api.OniPixelFormat.ONI_PIXEL_FORMAT_DEPTH_1_MM, resolutionX=w, resolutionY=h,
fps=fps))
## Check and configure the mirroring -- default is True
## Note: I disable mirroring
# print 'Mirroring info1', depth_stream.get_mirroring_enabled()
depth_stream.set_mirroring_enabled(False)
rgb_stream.set_mirroring_enabled(False)
## Start the streams
rgb_stream.start()
depth_stream.start()
## Synchronize the streams
dev.set_depth_color_sync_enabled(True) # synchronize the streams
## IMPORTANT: ALIGN DEPTH2RGB (depth wrapped to match rgb stream)
dev.set_image_registration_mode(openni2.IMAGE_REGISTRATION_DEPTH_TO_COLOR)
saving_folder_path = './shapenetcore_partanno_segmentation_benchmark_v0/tools/'
if not os.path.exists(saving_folder_path):
os.makedirs(saving_folder_path+'RGB')
os.makedirs(saving_folder_path+'D')
os.makedirs(saving_folder_path+'PC')
os.makedirs(saving_folder_path+'points')
os.makedirs(saving_folder_path+'points_label')
from config import CAMERA_CONFIG
## main loop
s = 1000
done = False
while not done:
key = cv2.waitKey(1) & 255
## Read keystrokes
if key == 27: # terminate
print("\tESC key detected!")
done = True
elif chr(key) == 's': # screen capture
print("\ts key detected. Saving image {}".format(s))
rgb = rgb[60:180, 80:240, :]
dmap = dmap[60:180, 80:240]
ply_content, points_content = generate_ply_from_rgbd(rgb=rgb, depth=dmap, config=CAMERA_CONFIG)
cv2.imwrite(saving_folder_path + "RGB/" + str(s) + '.png', rgb)
cv2.imwrite(saving_folder_path + "D/" + str(s) + '.png', dmap)
print(rgb.shape, dmap.shape)
print(type(rgb), type(dmap))
with open(saving_folder_path + "PC/" + str(s) + '.ply', 'w') as output:
output.write(ply_content)
print(saving_folder_path + "PC/" + str(s) + '.ply', ' done')
s += 1 # uncomment for multiple captures
# ### Get pointcloud of scene for prediction
# points_np = (np.array(points_content)[:, :3]).astype(np.float32)
# choice = np.random.choice(len(points_np), num_points, replace=True)
# points_np = points_np[choice, :]
# points_torch = torch.from_numpy(points_np)
#
# points_torch = points_torch.transpose(1, 0).contiguous()
#
# points_torch = Variable(points_torch.view(1, points_torch.size()[0], points_torch.size()[1]))
#
# ### Predict to segment scene
# pred, _ = classifier(points_torch)
# pred_choice = pred.data.max(2)[1]
# print(pred_choice)
## Streams
# RGB
rgb = get_rgb(rgb_stream=rgb_stream, h=h, w=w)
# DEPTH
dmap, d4d = get_depth(depth_stream=depth_stream, h=h, w=w)
# canvas
canvas = np.hstack((rgb, d4d))
## Display the stream syde-by-side
cv2.imshow('depth || rgb', canvas)
# end while
## Release resources
cv2.destroyAllWindows()
rgb_stream.stop()
depth_stream.stop()
openni2.unload()
print("Terminated")
loss.backward()
optimizer.step()
pred_choice = pred.data.max(1)[1]
correct = pred_choice.eq(target.data).cpu().sum()
print('[%d: %d/%d] train loss: %f accuracy: %f' %
(epoch, i, num_batch, loss.item(),
correct.item() / float(opt.batchSize * opt.num_points)))
if i % 100 == 0:
j, data = next(enumerate(testdataloader, 0))
points, target = data
points, target = Variable(points), Variable(target)
points = points.transpose(2, 1)
points, target = points.cuda(), target.cuda()
classifier = classifier.eval()
pred, _ = classifier(points)
pred = pred.view(-1, num_classes)
target = target.view(-1, 1)[:, 0] - 1
loss = F.nll_loss(pred, target)
pred_choice = pred.data.max(1)[1]
correct = pred_choice.eq(target.data).cpu().sum()
ioumax = func_miou(num_classes, target, pred_choice)
iou = sum(ioumax) / len(ioumax)
miou_list.append(iou)
miou = np.mean(miou_list)
print('[%d: %d/%d] %s loss: %f accuracy: %f IOU: %f mIOU %f' %
(epoch, i, num_batch, blue('test'), loss.item(),
correct.item() / float(opt.batchSize * opt.num_points), iou,
points, target = points.cuda(), target.cuda()
optimizer.zero_grad()
classifier = classifier.train()
pred, _ = classifier(points)
pred = pred.view(-1, num_classes)
target = target.view(-1,1)[:,0] - 1
#print(pred.size(), target.size())
loss = F.nll_loss(pred, target)
loss.backward()
optimizer.step()
pred_choice = pred.data.max(1)[1]
correct = pred_choice.eq(target.data).cpu().sum()
print('[%d: %d/%d] train loss: %f accuracy: %f' %(epoch, i, num_batch, loss.item(), correct.item()/float(opt.batchSize * 2500)))
if i % 10 == 0:
j, data = next(enumerate(testdataloader, 0))
points, target = data
points, target = Variable(points), Variable(target)
points = points.transpose(2,1)
points, target = points.cuda(), target.cuda()
classifier = classifier.eval()
pred, _ = classifier(points)
pred = pred.view(-1, num_classes)
target = target.view(-1,1)[:,0] - 1
loss = F.nll_loss(pred, target)
pred_choice = pred.data.max(1)[1]
correct = pred_choice.eq(target.data).cpu().sum()
print('[%d: %d/%d] %s loss: %f accuracy: %f' %(epoch, i, num_batch, blue('test'), loss.item(), correct.item()/float(opt.batchSize * 2500)))
torch.save(classifier.state_dict(), '%s/seg_model_%d.pth' % (opt.outf, epoch))
mt_classifier.load_state_dict(torch.load(opt.mt_model))
print('loading segmentation network for non-damaged data')
seg_classifier_all = PointNetDenseCls(k=dataset.num_seg_classes)
seg_classifier_all.load_state_dict(torch.load(opt.seg_all_model))
print('loading classification network for non-damaged data')
cls_classifier_all = PointNetCls(k=len(dataset.classes))
cls_classifier_all.load_state_dict(torch.load(opt.cls_all_model))
print('loading multi-task network for non-damaged data')
mt_classifier_all = PointNetMultiTask(cls_k=len(dataset.classes),
seg_k=dataset.num_seg_classes)
mt_classifier_all.load_state_dict(torch.load(opt.mt_all_model))
seg_classifier.eval()
cls_classifier.eval()
mt_classifier.eval()
seg_classifier_all.eval()
cls_classifier_all.eval()
mt_classifier_all.eval()
point = point.transpose(1, 0).contiguous()
point = point.view(1, point.size()[0], point.size()[1])
pred_cls_all, _ = cls_classifier_all(point)
pred_cls_all_choice = pred_cls_all.data.max(1)[1]
pred_seg_all, _ = seg_classifier_all(point)
pred_seg_all_choice = pred_seg_all.data.max(2)[1]
pred_cls_all_mt, preg_seg_all_mt, _ = mt_classifier_all(point)
pred_cls_mt_all_choice = pred_cls_all_mt.data.max(1)[1]
pred_seg_mt_all_choice = preg_seg_all_mt.data.max(2)[1]
print("model %d/%d" %( idx, len(d)))
point, seg = d[idx]
print(point.size(), seg.size())
point_np = point.numpy()
cmap = plt.cm.get_cmap("hsv", 10)
cmap = np.array([cmap(i) for i in range(10)])[:,:3]
gt = cmap[seg.numpy() - 1, :]
classifier = PointNetDenseCls(k = 4)
classifier.load_state_dict(torch.load(opt.model))
classifier.eval()
point = point.transpose(1,0).contiguous()
point = Variable(point.view(1, point.size()[0], point.size()[1]))
pred, _ = classifier(point)
pred_choice = pred.data.max(2)[1]
print(pred_choice)
#print(pred_choice.size())
pred_color = cmap[pred_choice.numpy()[0], :]
#print(pred_color.shape)
showpoints(point_np, gt, pred_color)
