# load the train, val and test dataset and pass to different functions
print('Loading required Tensors....\n')
train_images = sorted(os.listdir(train_folder_path))
total_images = len(train_images)
X_train = torch.zeros(total_images, num_channels, height, width)
# Y_train = torch.tensor([0] * total_images)
Y_train = torch.tensor([0] * total_images)
bounding_box = None
if train_method == 'bbox' or train_method == 'blackout':
# load bounding box information for each training image
bounding_box = {}
bbox = BoundingBox(train_folder_path, images_text_file, bounding_box_file, height, width)
for idx, image_name in enumerate(train_images):
bounding_box[idx] = bbox.get_bbox(image_name)
# if train with blackout - black out the extra region while loading
if train_method == 'blackout':
for idx, img_name in enumerate(train_images):
img_path = train_folder_path + '/' + img_name
img = Image.open(img_path)
if img.mode == 'L':
img = img.convert('RGB')
img = img.resize((width, height))
img = np.array(img)
x1, y1, x2, y2 = bounding_box[idx]
tmp = np.zeros((height, width, num_channels), dtype=np.uint8)
tmp[y1:y2, x1:x2, :] = 1
img = img * tmp
