diff --git a/data.py b/data.py
index 2406026381b75ebd85bb3194f92f0135aecb2812..e23af56f920968a718ed492e96a4a75d2b6b2121 100644
--- a/data.py
+++ b/data.py
@@ -73,11 +73,16 @@ def getMasterList_Fruits(basePath): # returns list with image, mask, and label
return [[a, b, c] for a, b, c in zip(imageList, maskList, labelList)]
-def classTrainingGenerator(model, batchSize, masterList=None, height=480, width=640, augmentation=True,
+
+def classTrainingGenerator(model, batchSize, masterList=None, height=480, width=640, augmentation=True, val=False,
**unused): # take input image, resize and store as rgb, create mask training data
- basePath = os.path.dirname(os.path.realpath(__file__)) + '/LINEMOD/' + model
+ basePath = os.path.dirname(os.path.realpath(__file__)) + '/Generated_Worlds_/Generated_Worlds_Training/' + model
+ if val:
+ basePath = os.path.dirname(
+ os.path.realpath(__file__)) + '/Generated_Worlds_/Generated_Worlds_Evaluating/' + model
+
if masterList is None:
- masterList = getMasterList(basePath)
+ masterList = getMasterList_Fruits(basePath)
random.shuffle(masterList)
i = 0
while True:
@@ -87,10 +92,10 @@ def classTrainingGenerator(model, batchSize, masterList=None, height=480, width=
if i == len(masterList):
i = 0
random.shuffle(masterList)
- x = filePathToArray(basePath + '/rgb/' + masterList[i][0], height, width)
+ x = filePathToArray(basePath + '/RGB_resized/' + masterList[i][0], height, width)
yClassLabels = np.zeros((height, width, 1)) # 1 class confidence value per model
- modelMask = filePathToArray(basePath + '/mask/' + masterList[i][1], height, width)
+ modelMask = filePathToArray(basePath + '/Instance_Mask_resized/' + masterList[i][1], height, width)
if augmentation:
if random.choice([True, False]): # vertical flip
@@ -112,13 +117,17 @@ def classTrainingGenerator(model, batchSize, masterList=None, height=480, width=
yield np.array(xBatch), np.array(yClassBatch)
-def coordsTrainingGenerator(model, batchSize, masterList=None, height=480, width=640, augmentation=True,
+def coordsTrainingGenerator(model, batchSize, masterList=None, val=False, height=480, width=640, augmentation=True,
altLabels=True): # takes input image and generates unit vector training data
print(f"-------- {batchSize}----------")
- basePath = os.path.dirname(os.path.realpath(__file__)) + '/LINEMOD/' + model
+ basePath = os.path.dirname(os.path.realpath(__file__)) + '/Generated_Worlds_/Generated_Worlds_Training/' + model
+ if val:
+ basePath = os.path.dirname(
+ os.path.realpath(__file__)) + '/Generated_Worlds_/Generated_Worlds_Evaluating/' + model
+
if masterList == None:
- masterList = getMasterList(basePath)
+ masterList = getMasterList_Fruits(basePath)
random.shuffle(masterList)
i = 0
while True:
@@ -128,14 +137,14 @@ def coordsTrainingGenerator(model, batchSize, masterList=None, height=480, width
if i == len(masterList):
i = 0
random.shuffle(masterList)
- x = filePathToArray(basePath + '/rgb/' + masterList[i][0], height, width)
+ x = filePathToArray(basePath + '/RGB_resized/' + masterList[i][0], height, width)
- with open(basePath + ('/altLabels/' if altLabels else '/labels/') + masterList[i][2]) as f:
+ with open(basePath + ('/altLabels/' if altLabels else '/FPS_resized/') + masterList[i][2]) as f:
labels = f.readline().split(' ')[1:19]
yCoordsLabels = np.zeros((height, width, 18)) # 9 coordinates
- modelMask = filePathToArray(basePath + '/mask/' + masterList[i][1], height, width)
+ modelMask = filePathToArray(basePath + '/Instance_Mask_resized/' + masterList[i][1], height, width)
if augmentation:
if random.choice([True, False]): # vertical flip
@@ -158,12 +167,16 @@ def coordsTrainingGenerator(model, batchSize, masterList=None, height=480, width
yield np.array(xBatch), np.array(yCoordBatch)
-def combinedTrainingGenerator(model, batchSize, masterList=None, height=480, width=640, out0='activation_9',
+def combinedTrainingGenerator(model, batchSize, masterList=None, val=False, height=480, width=640, out0='activation_9',
out1='activation_10', augmentation=True,
altLabels=True): # take input image, resize and store as rgb, create training data
- basePath = os.path.dirname(os.path.realpath(__file__)) + '/LINEMOD/' + model
+ basePath = os.path.dirname(os.path.realpath(__file__)) + '/Generated_Worlds_/Generated_Worlds_Training/' + model
+ if val:
+ basePath = os.path.dirname(
+ os.path.realpath(__file__)) + '/Generated_Worlds_/Generated_Worlds_Evaluating/' + model
+
if masterList is None:
- masterList = getMasterList(basePath)
+ masterList = getMasterList_Fruits(basePath)
i = 0
while True:
xBatch = []
@@ -173,16 +186,16 @@ def combinedTrainingGenerator(model, batchSize, masterList=None, height=480, wid
if i == len(masterList):
i = 0
random.shuffle(masterList)
- x = filePathToArray(basePath + '/rgb/' + masterList[i][0], height, width)
+ x = filePathToArray(basePath + '/RGB_resized/' + masterList[i][0], height, width)
- with open(basePath + ('/altLabels/' if altLabels else '/labels/') + masterList[i][2]) as f:
+ with open(basePath + ('/altLabels/' if altLabels else '/FPS_resized/') + masterList[i][2]) as f:
labels = f.readline().split(' ')[1:19]
yCoordsLabels = np.zeros((height, width, 18)) # 9 coordinates
yClassLabels = np.zeros((height, width, 1)) # 1 class confidence value per model
# yClassLabels = np.tile(np.array([1, 0]),(height, width, 1))
- modelMask = filePathToArray(basePath + '/mask/' + masterList[i][1], height, width)
+ modelMask = filePathToArray(basePath + '/Instance_Mask_resized/' + masterList[i][1], height, width)
if augmentation:
# for data aug, get random horizontal, vertical flips, flip input x with np, label vals = 1 - labelvals, flip mask
@@ -293,17 +306,18 @@ def getDataSplit(genNew=False, split=.8, modelClass='cat'):
splitDict = pickle.load(f)
return splitDict["trainData"], splitDict["validData"]
-def getDataSplit_Fruits(genNew=False, split=.8, modelClass='cat'):
+
+def getDataSplit_Fruits(genNew=True, split=.8, modelClass='cat'):
# access training data, get jpeg, mask, label filenames split into training / validation sets
if genNew: # create split
- basePathTraining = f'/home/gduret/Documents/guimod/Generated_Worlds_Training/{modelClass}' # os.path.dirname(os.path.realpath(__file__)) + '/LINEMOD/' + modelClass
- basePathEvaluating = f'/home/gduret/Documents/guimod/Generated_Worlds_Evaluating/{modelClass}'
+ basePathTraining = f'/home/mahmoud/pvnet_method/Generated_Worlds_/Generated_Worlds_Training/{modelClass}' # os.path.dirname(os.path.realpath(__file__)) + '/LINEMOD/' + modelClass
+ basePathEvaluating = f'/home/mahmoud/pvnet_method/Generated_Worlds_/Generated_Worlds_Evaluating/{modelClass}'
masterList_Training = getMasterList_Fruits(basePathTraining)
masterList_Evalution = getMasterList_Fruits(basePathEvaluating)
random.shuffle(masterList_Training)
random.shuffle(masterList_Evalution)
- #splitPoint = round(len(masterList) * split)
+ # splitPoint = round(len(masterList) * split)
splitDict = {"trainData": masterList_Training, "validData": masterList_Evalution}
@@ -316,7 +330,6 @@ def getDataSplit_Fruits(genNew=False, split=.8, modelClass='cat'):
return splitDict["trainData"], splitDict["validData"]
-
def genAltLabels(p3dOld, p3dNew, matrix=np.array([[572.4114, 0., 325.2611], [0., 573.57043, 242.04899], [0., 0., 1.]]),
method=cv2.SOLVEPNP_ITERATIVE, modelClass='cat', height=480, width=640,
showPoint=False): # generate pixel labels for p3dNew using labels for p3dOld