def test_rot90(target):
img = np.array([[0, 0, 1], [0, 0, 1], [0, 0, 1]], dtype=np.uint8)
expected = np.array([[1, 1, 1], [0, 0, 0], [0, 0, 0]], dtype=np.uint8)
if target == 'image':
img = convert_2d_to_3d(img)
expected = convert_2d_to_3d(expected)
rotated = F.rot90(img, factor=1)
assert np.array_equal(rotated, expected)
def test_rot90_float(target):
img = np.array([[0.0, 0.0, 0.4], [0.0, 0.0, 0.4], [0.0, 0.0, 0.4]],
dtype=np.float32)
expected = np.array([[0.4, 0.4, 0.4], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
dtype=np.float32)
img, expected = convert_2d_to_target_format([img, expected], target=target)
rotated = F.rot90(img, factor=1)
assert_array_almost_equal_nulp(rotated, expected)
def test_rot90(target):
img = np.array(
[[0, 0, 1],
[0, 0, 1],
[0, 0, 1]], dtype=np.uint8)
expected = np.array(
[[1, 1, 1],
[0, 0, 0],
[0, 0, 0]], dtype=np.uint8)
img, expected = convert_2d_to_target_format([img, expected], target=target)
rotated = F.rot90(img, factor=1)
assert np.array_equal(rotated, expected)
def test_rot90(target):
img = np.array(
[[0, 0, 1],
[0, 0, 1],
[0, 0, 1]], dtype=np.uint8)
expected = np.array(
[[1, 1, 1],
[0, 0, 0],
[0, 0, 0]], dtype=np.uint8)
if target == 'image':
img = convert_2d_to_3d(img)
expected = convert_2d_to_3d(expected)
rotated = F.rot90(img, factor=1)
assert np.array_equal(rotated, expected)
def __getitem__(self, index: int):
while True:
video, img_file, label, ori_video, frame, fold = self.data[index]
try:
if self.mode == "train":
label = np.clip(label, self.label_smoothing,
1 - self.label_smoothing)
img_path = os.path.join(self.data_root, self.crops_dir, video,
img_file)
image = cv2.imread(img_path, cv2.IMREAD_COLOR)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_width, image_height = image.shape[:2]
mask = np.zeros(image.shape[:2], dtype=np.uint8)
diff_path = os.path.join(self.data_root, "diffs", video,
img_file[:-4] + "_diff.png")
landmark_path = os.path.join(self.data_root, "landmarks",
ori_video, img_file[:-4] + ".npy")
try:
msk = cv2.imread(diff_path, cv2.IMREAD_GRAYSCALE)
if msk is not None:
mask = msk
except:
print("not found mask", diff_path)
pass
if self.mode == "train" and self.hardcore and not self.rotation:
if os.path.exists(landmark_path) and random.random() < 0.7:
landmarks = np.load(landmark_path)
image = remove_landmark(image, landmarks)
elif random.random() < 0.2:
blackout_convex_hull(image)
elif random.random() < 0.1:
binary_mask = mask > 0.4 * 255
masks = prepare_bit_masks(
(binary_mask * 1).astype(np.uint8))
tries = 6
current_try = 1
while current_try < tries:
bitmap_msk = random.choice(masks)
if label < 0.5 or np.count_nonzero(
mask * bitmap_msk) > 20:
mask *= bitmap_msk
image *= np.expand_dims(bitmap_msk, axis=-1)
break
current_try += 1
if self.mode == "train" and self.padding_part > 3:
image = change_padding(image, self.padding_part)
valid_label = np.count_nonzero(
mask[mask > 20]) > 32 or label < 0.5
valid_label = 1 if valid_label else 0
rotation = 0
if self.transforms:
data = self.transforms(image=image, mask=mask)
image = data["image"]
mask = data["mask"]
if self.mode == "train" and self.hardcore and self.rotation:
# landmark_path = os.path.join(self.data_root, "landmarks", ori_video, img_file[:-4] + ".npy")
dropout = 0.8 if label > 0.5 else 0.6
if self.rotation:
dropout *= 0.7
elif random.random() < dropout:
blackout_random(image, mask, label)
#
# os.makedirs("../images", exist_ok=True)
# cv2.imwrite(os.path.join("../images", video+ "_" + str(1 if label > 0.5 else 0) + "_"+img_file), image[...,::-1])
if self.mode == "train" and self.rotation:
rotation = random.randint(0, 3)
image = rot90(image, rotation)
image = img_to_tensor(image, self.normalize)
# assuming the image size is 380*380
image_size = 380
image = Resize(image_size, image_size)
rect = detector(img)[0]
x = rect.left()
y = rect.top()
w = rect.right() - x
h = rect.buttom() - y
top_left = np.array([x, y])
top_right = np.array([x + w, y])
bottom_left = np.array([x, y + h])
bottom_right = np.array([x + w, y + h])
# insert the face rectangle points into the landmark data
landmark = np.insert(landmark, 0, top_left, axis=0)
landmark = np.insert(landmark, 0, top_right, axis=0)
landmark = np.insert(landmark, 0, bottom_left, axis=0)
landmark = np.insert(landmark, 0, bottom_right, axis=0)
# perform scaling for the landmark
for i in landmark:
i[0] = round(image_size * (i[0] / image_width))
i[1] = round(image_size * (i[1] / image_height))
return {
"image": image,
"labels": np.array((label, )),
"img_name": os.path.join(video, img_file),
"valid": valid_label,
"rotations": rotation,
"landmark": landmark
}
except Exception as e:
traceback.print_exc(file=sys.stdout)
print(
"Broken image",
os.path.join(self.data_root, self.crops_dir, video,
img_file))
index = random.randint(0, len(self.data) - 1)
def __getitem__(self, index: int):
while True:
video, img_file, label, ori_video, frame, fold = self.data[index]
try:
if self.mode == "train":
label = np.clip(label, self.label_smoothing,
1 - self.label_smoothing)
img_path = os.path.join(self.data_root, self.crops_dir, video,
img_file)
print("img path plis %s" % img_path)
image = cv2.imread(img_path, cv2.IMREAD_COLOR)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
mask = np.zeros(image.shape[:2], dtype=np.uint8)
diff_path = os.path.join(self.data_root, "diffs", video,
img_file[:-4] + "_diff.png")
print("diff path plis %s" % diff_path)
try:
msk = cv2.imread(diff_path, cv2.IMREAD_GRAYSCALE)
if msk is not None:
mask = msk
except:
print("not found mask", diff_path)
pass
if self.mode == "train" and self.hardcore and not self.rotation:
landmark_path = os.path.join(self.data_root, "landmarks",
ori_video,
img_file[:-4] + ".npy")
if os.path.exists(landmark_path) and random.random() < 0.7:
landmarks = np.load(landmark_path)
image = remove_landmark(image, landmarks)
elif random.random() < 0.2:
blackout_convex_hull(image)
elif random.random() < 0.1:
binary_mask = mask > 0.4 * 255
masks = prepare_bit_masks(
(binary_mask * 1).astype(np.uint8))
tries = 6
current_try = 1
while current_try < tries:
bitmap_msk = random.choice(masks)
if label < 0.5 or np.count_nonzero(
mask * bitmap_msk) > 20:
mask *= bitmap_msk
image *= np.expand_dims(bitmap_msk, axis=-1)
break
current_try += 1
if self.mode == "train" and self.padding_part > 3:
image = change_padding(image, self.padding_part)
valid_label = np.count_nonzero(
mask[mask > 20]) > 32 or label < 0.5
valid_label = 1 if valid_label else 0
print("tipo valid label %d" % valid_label)
rotation = 0
if self.transforms:
data = self.transforms(image=image, mask=mask)
image = data["image"]
mask = data["mask"]
print("nu transformed kazka")
if self.mode == "train" and self.hardcore and self.rotation:
# landmark_path = os.path.join(self.data_root, "landmarks", ori_video, img_file[:-4] + ".npy")
dropout = 0.8 if label > 0.5 else 0.6
if self.rotation:
dropout *= 0.7
elif random.random() < dropout:
blackout_random(image, mask, label)
#
# os.makedirs("../images", exist_ok=True)
# cv2.imwrite(os.path.join("../images", video+ "_" + str(1 if label > 0.5 else 0) + "_"+img_file), image[...,::-1])
if self.mode == "train" and self.rotation:
rotation = random.randint(0, 3)
image = rot90(image, rotation)
image = img_to_tensor(image, self.normalize)
print("nu i tensor daba")
return {
"image": image,
"labels": np.array((label, )),
"img_name": os.path.join(video, img_file),
"valid": valid_label,
"rotations": rotation
}
except Exception as e:
print("nu tiesiog exception %s" % sys.exc_info()[0])
traceback.print_exc(file=sys.stdout)
print(
"Broken image",
os.path.join(self.data_root, self.crops_dir, video,
img_file))
index = random.randint(0, len(self.data) - 1)
def __getitem__(self, index: int):
while True:
video, img_file, label, ori_video, frame, fold = self.data[index]
try:
if self.mode == "train":
label = np.clip(label, self.label_smoothing,
1 - self.label_smoothing)
img_path = os.path.join(self.data_root, self.crops_dir, video,
img_file)
# get the image from the directory
image = cv2.imread(img_path, cv2.IMREAD_COLOR)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# create a mask with the image
mask = np.zeros(image.shape[:2], dtype=np.uint8)
# gets the diffrences image (SSIM image from the preprocessing)
# of the original image and fake image if it exists
diff_path = os.path.join(self.data_root, "diffs", video,
img_file[:-4] + "_diff.png")
try:
msk = cv2.imread(diff_path, cv2.IMREAD_GRAYSCALE)
if msk is not None:
mask = msk
except:
print("not found mask", diff_path)
pass
if self.mode == "train" and self.hardcore and not self.rotation:
# this part of the function gets the image and uses random() to choose
# one of the following generalization approaches randomly:
# 1. randomly removes one of the landmarks that MTCNN detected on the face randomly:
# it blacks out landmarks (eyes, nose or mouth)
# 2. removes/blacks out half face horisontally or vertically and it uses dlib face convex hulls to do that
# 3. blacks out half the image
landmark_path = os.path.join(self.data_root, "landmarks",
ori_video,
img_file[:-4] + ".npy")
if os.path.exists(landmark_path) and random.random() < 0.7:
landmarks = np.load(landmark_path)
image = remove_landmark(image, landmarks)
elif random.random() < 0.2:
blackout_convex_hull(image)
elif random.random() < 0.1:
binary_mask = mask > 0.4 * 255
masks = prepare_bit_masks(
(binary_mask * 1).astype(np.uint8))
tries = 6
current_try = 1
while current_try < tries:
bitmap_msk = random.choice(masks)
if label < 0.5 or np.count_nonzero(
mask * bitmap_msk) > 20:
mask *= bitmap_msk
image *= np.expand_dims(bitmap_msk, axis=-1)
break
current_try += 1
if self.mode == "train" and self.padding_part > 3:
image = change_padding(image, self.padding_part)
valid_label = np.count_nonzero(
mask[mask > 20]) > 32 or label < 0.5
valid_label = 1 if valid_label else 0
rotation = 0
if self.transforms:
data = self.transforms(image=image, mask=mask)
image = data["image"]
mask = data["mask"]
if self.mode == "train" and self.hardcore and self.rotation:
# landmark_path = os.path.join(self.data_root, "landmarks", ori_video, img_file[:-4] + ".npy")
dropout = 0.8 if label > 0.5 else 0.6
if self.rotation:
dropout *= 0.7
elif random.random() < dropout:
blackout_random(image, mask, label)
# os.makedirs("../images", exist_ok=True)
# cv2.imwrite(os.path.join("../images", video+ "_" + str(1 if label > 0.5 else 0) + "_"+img_file), image[...,::-1])
if self.mode == "train" and self.rotation:
rotation = random.randint(0, 3)
image = rot90(image, rotation)
image = img_to_tensor(image, self.normalize)
return {
"image": image,
"labels": np.array((label, )),
"img_name": os.path.join(video, img_file),
"valid": valid_label,
"rotations": rotation
}
except Exception as e:
traceback.print_exc(file=sys.stdout)
print(
"Broken image",
os.path.join(self.data_root, self.crops_dir, video,
img_file))
index = random.randint(0, len(self.data) - 1)
def __getitem__(self, index: int):
# while(True):
video, img_file, label, ori_video, frame, fold = self.data[index]
# try:
if self.mode == "train":
label = np.clip(label, self.label_smoothing,
1 - self.label_smoothing)
# Load image and mask
img_path = os.path.join(self.data_root, self.crops_dir, video,
img_file)
image = cv2.imread(img_path, cv2.IMREAD_COLOR)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# Applying hardcore augmentations without rotation
if self.mode == "train" and self.hardcore and not self.rotation and not self.random_erase:
landmark_path = os.path.join(self.data_root, "landmarks",
ori_video, img_file[:-4] + ".npy")
# Remove facial features using landmark informations
if os.path.exists(landmark_path) and random.random() < 0.7:
landmarks = np.load(landmark_path)
image = remove_landmark(image, landmarks, self.cutout_fill)
# Remove facial parts using convex hull
elif random.random() < 0.4:
err = 0
cp = np.copy(image)
try:
blackout_convex_hull(cp, detector, predictor,
self.cutout_fill)
except Exception:
err = 1
if err == 0:
image = cp
# Remove parts of image randomly from 6 bitmasks
# elif random.random() < 0.1:
# binary_mask = mask > 0.4 * 255
# masks = prepare_bit_masks((binary_mask * 1).astype(np.uint8))
# tries = 6
# current_try = 1
# while current_try < tries:
# bitmap_msk = random.choice(masks)
# if label < 0.5 or np.count_nonzero(mask * bitmap_msk) > 20:
# mask *= bitmap_msk
# image *= np.expand_dims(bitmap_msk, axis=-1)
# break
# current_try += 1
# Resize image and remove extra padding outside face
if self.mode == "train" and self.padding_part > 3:
image = change_padding(image, self.padding_part)
# Change label depending on ratio of fake parts in mask
# TODO : change mask to binary_mask
# valid_label = np.count_nonzero(mask[mask > 20]) > 32 or label < 0.5
# valid_label = 1 if valid_label else 0
# Use builtin transforms passed in
if self.transforms is not None:
data = self.transforms(image=image)
image = data["image"]
# mask = data["mask"]
# If hardcore augmentation and rotation are both true
# then applies only random blackout
# if self.mode == "train" and self.hardcore and self.rotation:
# dropout = 0.8 if label > 0.5 else 0.6
# if self.rotation:
# dropout *= 0.7
# elif random.random() < dropout:
# blackout_random(image, mask, label)
rotation = 0
# Applies 90 degree rotation
if self.mode == "train" and self.rotation:
rotation = random.randint(0, 3)
image = rot90(image, rotation)
# Saves 20% of the train images
# if(random.random() < 0.1 and conv == True):
# os.makedirs("train_images", exist_ok=True)
# cv2.imwrite(os.path.join("train_images", video+ "_" + str(1 if label > 0.5 else 0) + "_"+img_file), image[...,::-1])
image = img_to_tensor(image, self.normalize)
if self.mode == "train" and self.random_erase:
image = RandomErasing(p=0.5, scale=(0.02, 0.2),
value="random")(image)
return {
"image": image,
"label": np.array((label, )),
"image_name": os.path.join(video, img_file),
"rotations": rotation
}
