im_height = 480
# Specify which gpus to use
torch.manual_seed(seed)
if use_cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = gpus
torch.cuda.manual_seed(seed)
# Specifiy the model and the loss
model = Darknet(cfgfile)
region_loss = model.loss
# Model settings
# model.load_weights(weightfile)
model.load_weights_until_last(weightfile)
model.print_network()
model.seen = 0
region_loss.iter = model.iter
region_loss.seen = model.seen
processed_batches = model.seen // batch_size
init_width = model.width
init_height = model.height
init_epoch = model.seen // nsamples
# Variable to save
training_iters = []
training_losses = []
testing_iters = []
testing_errors_pixel = []
testing_accuracies = []
def test(datacfg, cfgfile, weightfile, imgfile):
# ******************************************#
# PARAMETERS PREPARATION #
# ******************************************#
#parse configuration files
options = read_data_cfg(datacfg)
meshname = options['mesh']
name = options['name']
#Parameters for the network
seed = int(time.time())
gpus = '0' # define gpus to use
test_width = 544 # define test image size
test_height = 544
torch.manual_seed(seed) # seed torch random
use_cuda = True
if use_cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = gpus
torch.cuda.manual_seed(seed) # seed cuda random
conf_thresh = 0.1
# Read object 3D model, get 3D Bounding box corners
mesh = MeshPly(meshname)
vertices = np.c_[np.array(mesh.vertices), np.ones((len(mesh.vertices), 1))].transpose()
corners3D = get_3D_corners(vertices)
diam = float(options['diam'])
# now configure camera intrinsics
internal_calibration = get_camera_intrinsic()
# ******************************************#
# NETWORK CREATION #
# ******************************************#
# Create the network based on cfg file
model = Darknet(cfgfile)
model.print_network()
model.load_weights(weightfile)
# Pass the model to GPU
if use_cuda:
# model = model.cuda()
model = torch.nn.DataParallel(model, device_ids=[0]).cuda() # Multiple GPU parallelism
model.eval()
num_classes = model.module.num_classes
anchors = model.module.anchors
num_anchors = model.module.num_anchors
# ******************************************#
# INPUT IMAGE PREPARATION FOR NN #
# ******************************************#
# Now prepare image: convert to RGB, resize, transform to Tensor
# use cuda,
img = Image.open(imgfile).convert('RGB')
ori_size = img.size # store original size
img = img.resize((test_width, test_height))
t1 = time.time()
img = transforms.Compose([transforms.ToTensor(),])(img)#.float()
img = Variable(img, requires_grad = True)
img = img.unsqueeze(0) # add a fake batch dimension
img = img.cuda()
# ******************************************#
# PASS IT TO NETWORK AND GET PREDICTION #
# ******************************************#
# Forward pass
output = model(img).data
#print("Output Size: {}".format(output.size(0)))
t2 = time.time()
# Reload Original img
img = cv2.imread(imgfile)
for k in range(num_classes):
all_boxes = get_corresponding_region_boxes(output, conf_thresh, num_classes, anchors, num_anchors, k, only_objectness=0)
t4 = time.time()
for i in range(output.size(0)):
# For each image, get all the predictions
boxes = all_boxes[i]
best_conf_est = -1
# If the prediction has the highest confidence, choose it as our prediction
for j in range(len(boxes)):
if (boxes[j][18] > best_conf_est) and (boxes[j][20] == k):
best_conf_est = boxes[j][18]
box_pr = boxes[j]
#bb2d_gt = get_2d_bb(box_gt[:18], output.size(3))
bb2d_pr = get_2d_bb(box_pr[:18], output.size(3))
#for a,b in zip(bb2d_gt, bb2d_pr):
# print(type(a),type(b))
#iou = bbox_iou(bb2d_gt, bb2d_pr)
#match = corner_confidence9(box_gt[:18], torch.FloatTensor(boxes[j][:18]))
corners2D_pr = np.array(np.reshape(box_pr[:18], [9, 2]), dtype='float32')
corners2D_pr[:, 0] = corners2D_pr[:, 0] * ori_size[0] # Width
corners2D_pr[:, 1] = corners2D_pr[:, 1] * ori_size[1] # Heightt
t3 = time.time()
# draw each predicted 2D point
for v, (x,y) in enumerate(corners2D_pr):
# get colors to draw the lines
col1 = 28*v
col2 = 255 - (28*v)
col3 = np.random.randint(0,256)
cv2.circle(img, (x,y), 3, (col1,col2,col3), -1)
cv2.putText(img, str(v), (int(x) + 5, int(y) + 5),cv2.FONT_HERSHEY_SIMPLEX, 0.5, (col1, col2, col3), 1)
# Get each predicted point and the centroid
p1 = corners2D_pr[1]
p2 = corners2D_pr[2]
p3 = corners2D_pr[3]
p4 = corners2D_pr[4]
p5 = corners2D_pr[5]
p6 = corners2D_pr[6]
p7 = corners2D_pr[7]
p8 = corners2D_pr[8]
center = corners2D_pr[0]
# Draw cube lines around detected object
# draw front face
line_point = 3
cv2.line(img,(p1[0],p1[1]),(p2[0],p2[1]), (0,255,0),line_point)
cv2.line(img,(p2[0],p2[1]),(p4[0],p4[1]), (0,255,0),line_point)
cv2.line(img,(p4[0],p4[1]),(p3[0],p3[1]), (0,255,0),line_point)
cv2.line(img,(p3[0],p3[1]),(p1[0],p1[1]), (0,255,0),line_point)
# draw back face
cv2.line(img,(p5[0],p5[1]),(p6[0],p6[1]), (0,255,0),line_point)
cv2.line(img,(p7[0],p7[1]),(p8[0],p8[1]), (0,255,0),line_point)
cv2.line(img,(p6[0],p6[1]),(p8[0],p8[1]), (0,255,0),line_point)
cv2.line(img,(p5[0],p5[1]),(p7[0],p7[1]), (0,255,0),line_point)
# draw right face
cv2.line(img,(p2[0],p2[1]),(p6[0],p6[1]), (0,255,0),line_point)
cv2.line(img,(p1[0],p1[1]),(p5[0],p5[1]), (0,255,0),line_point)
# draw left face
cv2.line(img,(p3[0],p3[1]),(p7[0],p7[1]), (0,255,0),line_point)
cv2.line(img,(p4[0],p4[1]),(p8[0],p8[1]), (0,255,0),line_point)
# create a window to display image
wname = "Prediction"
cv2.namedWindow(wname)
# Show the image and wait key press
cv2.imshow(wname, img)
cv2.waitKey()
print(output.shape)
