def main():
# Basic set
patchsize = 40
# Load networks
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_obj = OBJ_CNN().to(DEVICE)
model_obj.load_state_dict(
torch.load('./Model parameter/fire/obj_model_init.pkl'))
# Traning path setting
dataDir = './'
trainingDir = dataDir + 'training/' # If u want to do experiment on test dataset, change it
trainingSets = util.getSubPaths(trainingDir)
# Initialize training datasets
trainingDataset = dataset.Dataset()
trainingDataset.readFileNames(trainingSets[0])
trainingDataset.SetObjID(1)
i = 523 # If u want to do experiment on other RGB image, change it
imgBGR = trainingDataset.getBGR(i)
# Do every pixel sample
sample = [[idx % 640, idx // 640] for idx in range(480 * 640)]
sample_with_patch = [idx for idx in sample if patchsize / 2 <= idx[0] < 640 - patchsize / 2 and \
patchsize / 2 <= idx[1] < 480 - patchsize / 2]
sample_with_patch = np.array(sample_with_patch).reshape([440, 600, 2])
pred_coord = cnn.getCoordImg(colorData=imgBGR,
sampling=sample_with_patch,
patchsize=patchsize,
model=model_obj) / 1000.0
return pred_coord
# Parameters setting
inputSize = 42
channels = 3
trainingLimit = 320000
trainingImages = 100
trainingPatches = 512
BATCHSIZE = 64
lrInitPre = 0.0001
storeIntervalPre = 1000
lrInterval = 50000
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Training path setting
dataDir = './'
trainingDir = dataDir + 'training/'
trainingSets = util.getSubPaths(trainingDir)
# Test path setting
dataDir = './'
TestDir = dataDir + 'test/'
TestSets = util.getSubPaths(TestDir)
# transform for data
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
'''
If we wanna use fire scene, change the following:
training.readFileNames(trainingSets[0]) -> training.readFileNames(trainingSets[1])
# Parameter setting
trainingImages = 100
trainingHyps = 16
trainingRounds = 80
objTemperature = 10
objBatchSize = 64
lrInitPre = 0.0001
CNN_OBJ_PATCHSIZE = 40
CNN_RGB_PATCHSIZE = 42
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Training path setting
dataDir = './'
trainingDir = dataDir + 'training/'
trainingSets = util.getSubPaths(trainingDir)
# Load RGB CNN's parameters
RGB_NET = OBJ_CNN()
RGB_NET.load_state_dict(torch.load('./Model parameter/obj_model_init.pkl'))
# Dataset process
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
])
# Initialize training dataset
trainingDataset = dataset.Dataset()
trainingDataset.readFileNames(trainingSets[0])
trainingDataset.SetObjID(1)
def main():
GlobalProp = properties.GlobalProperties()
objHyps = GlobalProp.ransacIterations
refSteps = GlobalProp.ransacRefinementIterations
inlierThreshold2D = GlobalProp.ransacInlierThreshold2D
refInlierCount = GlobalProp.ransacBatchSize
ptCount = 4
CNN_OBJ_PATCHSIZE = 40
camMat = GlobalProp.getCamMat()
CamMat = camMat.astype(np.float)
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Traning path setting
dataDir = './'
trainingDir = dataDir + 'training/'
trainingSets = util.getSubPaths(trainingDir)
print("Loading train set...\n")
# Initialize training datasets
training = dataset.Dataset()
training.readFileNames(trainingSets[0])
training.SetObjID(1)
# For loss recording
testfile = open('./Model parameter/chess/scene/ransac_test_loss.txt', 'a')
# set for coordinate net
model_obj = OBJ_CNN().to(DEVICE)
model_obj.load_state_dict(
torch.load('./Model parameter/chess/scene/obj_model_init.pkl'))
# set for score net
model_score = SCORE_CNN().to(DEVICE)
model_score.load_state_dict(
torch.load('./Model parameter/chess/scene/score_model_init.pkl'))
model_obj.eval()
model_score.eval()
avgCorrect = 0
expLosses = []
sfEntropies = []
rotErrs = []
tErrs = []
# For errors recording
testErrfile = open('./Model parameter/chess/scene/ransac_test_errors.txt',
'a')
# Training set
TrainingRound = 5000
round = 0
# For recording
# E2E_FILE = './end2end/'
# if not os.path.exists(E2E_FILE):
# os.mkdir(E2E_FILE)
# OBJ_E2E = E2E_FILE + 'model_obj_e2e.pkl'
# SCORE_E2E = E2E_FILE + 'model_score_e2e.pkl'
PARAMETER = list(model_obj.parameters()) + list(model_score.parameters())
OPTIMIZER = torch.optim.SGD(PARAMETER, lr=lrInitE2E, momentum=momentumE2E)
# OPTIMZER_OBJ = torch.optim.SGD(model_obj.parameters(), lr=lrInitE2E, momentum=momentumE2E)
# OPTIMZER_SCORE = torch.optim.SGD(model_score.parameters(), lr=lrInitE2E, momentum=momentumE2E)
time_start = time.time()
argcorrect = 0
argcorrect_score = 0
argcorrect_ransac = 0
dsac_errfile = open('./Model parameter/chess/scene/dsac_err.txt', 'a')
ransac_errfile = open('./Model parameter/chess/scene/ransac_err.txt', 'a')
inlier_errfile = open('./Model parameter/chess/scene/inlier_err.txt', 'a')
# Beginning
# len(testDataset)
for i in range(0, 2000):
print("Processing test image no.", i + 1)
testBGR = training.getBGR(i)
testInfo = training.getInfo(i)
# process frame (same function used in training, hence most of the variables below are not used here), see method documentation for parameter explanation
hyps = []
refHyps = []
# imgPts = []
# objPts = []
imgIdx = []
estObj = np.zeros((CNN_OBJ_PATCHSIZE, CNN_OBJ_PATCHSIZE, 3))
sampling = np.zeros((CNN_OBJ_PATCHSIZE, CNN_OBJ_PATCHSIZE, 3))
sampledPoints = []
inlierMaps = []
pixelIdxs = []
correct, correct_score, correct_ransac,\
rotErr, tErr, \
rotErr_score, tErr_score,\
rotErr_ransac, tErr_ransac = cnn.processImage(
testBGR,
testInfo,
model_obj,
model_score,
objHyps,
ptCount,
CamMat,
inlierThreshold2D,
refInlierCount,
refSteps,
refHyps,
sampledPoints,
estObj,
sampling,
inlierMaps,
pixelIdxs,
OPTIMIZER,
round)
time_end = time.time()
print('total cost time:', time_end - time_start)
if correct:
argcorrect += 1
print('argcorrect:', argcorrect)
if correct_score:
argcorrect_score += 1
print('argcorrect_score:', argcorrect_score)
if correct_ransac:
argcorrect_ransac += 1
print('argcorrect_ransac:', argcorrect_ransac)
np.savetxt(dsac_errfile, np.array([rotErr_score, tErr_score]))
np.savetxt(ransac_errfile, np.array([rotErr_ransac, tErr_ransac]))
np.savetxt(inlier_errfile, np.array([rotErr, tErr]))
dsac_errfile.close()
ransac_errfile.close()
inlier_errfile.close()