Commit 77ba43b1 authored by Trái Vú Sữa's avatar Trái Vú Sữa 😄

fix #2

parent 186646b8
......@@ -211,6 +211,7 @@ def _resize_subtract_mean(image, insize, rgb_mean):
image = cv2.resize(image, (insize, insize), interpolation=interp_method)
image = image.astype(np.float32)
image -= rgb_mean
image /= 128.0
return image.transpose(2, 0, 1)
......
import os
import os.path
import sys
import torch
import torch.utils.data as data
import cv2
import numpy as np
import torch.nn.functional as F
import skimage.transform
import torchvision.transforms as transforms
from imgaug import augmenters as iaa
class FaceDataset(data.Dataset):
......@@ -15,15 +12,15 @@ class FaceDataset(data.Dataset):
super(FaceDataset, self).__init__()
self.path_images, self.labels = self.read_file(root_path, file_name)
self.preproc = preproc
self.augment = ImgAugTransform()
self.target_transform = target_transform
def __len__(self):
return len(self.path_images)
def __getitem__(self, idx):
img = cv2.imread(self.path_images[idx])
img_raw = cv2.imread(self.path_images[idx])
img = self.augment(img_raw)
labels = self.labels[idx]
annotations = np.zeros((0, 15))
......@@ -55,7 +52,7 @@ class FaceDataset(data.Dataset):
annotations = np.append(annotations, annotation, axis=0)
target = np.array(annotations)
debug = False
debug = True
if debug:
img_debug = img.copy()
for index, b in enumerate(annotations):
......@@ -69,8 +66,9 @@ class FaceDataset(data.Dataset):
cv2.circle(img_debug, (b[10], b[11]), 1, (0, 255, 0), 4)
cv2.circle(img_debug, (b[12], b[13]), 1, (255, 0, 0), 4)
name = "test_data.jpg"
cv2.imwrite(name, img_debug)
img_show = np.hstack((img_debug, img_raw))
cv2.imshow("test", img_show)
cv2.waitKey()
if self.preproc is not None:
img, target = self.preproc(img, target)
......@@ -78,8 +76,8 @@ class FaceDataset(data.Dataset):
labels = target[:, -1]
landms = target[:, 4:14]
# TODO write landms to target_transforms
if self.target_transform:
boxes,landms, labels = self.target_transform(boxes,landms,labels)
# if self.target_transform:
# boxes,landms, labels = self.target_transform(boxes,landms,labels)
return torch.from_numpy(img), target
@staticmethod
......@@ -105,3 +103,101 @@ class FaceDataset(data.Dataset):
labels.append(label)
words.append(labels)
return path_images, words
class ImgAugTransform:
def __init__(self):
self.aug = iaa.Sequential([
iaa.Sometimes(0.15, iaa.Grayscale(alpha=(0.05, 1.0))),
iaa.Sometimes(0.2,
iaa.OneOf([
iaa.Sometimes(0.25, iaa.MotionBlur(k=5, angle=[-45, 45])),
iaa.Sometimes(0.25, iaa.GaussianBlur(sigma=(0, 1.3)))
])),
iaa.Sometimes(0.2,
iaa.OneOf([
iaa.Sharpen(alpha=(0.0, 1.0), lightness=(0.75, 2.0)),
iaa.Emboss(alpha=(0.0, 1.0), strength=(0.5, 1.5)),
iaa.EdgeDetect(alpha=(0.0, 1.0)),
iaa.DirectedEdgeDetect(alpha=(0.0, 1.0), direction=(0.0, 1.0)),
iaa.Canny(
alpha=(0.0, 0.5),
colorizer=iaa.RandomColorsBinaryImageColorizer(
color_true=255,
color_false=0
)
)
])),
iaa.Sometimes(0.2,
iaa.OneOf([
iaa.GammaContrast((0.5, 2.0)),
iaa.GammaContrast((0.5, 2.0), per_channel=True),
iaa.SigmoidContrast(gain=(3, 10), cutoff=(0.4, 0.6)),
iaa.SigmoidContrast(gain=(3, 10), cutoff=(0.4, 0.6), per_channel=True),
iaa.LogContrast(gain=(0.6, 1.4)),
iaa.LinearContrast((0.4, 1.6)),
iaa.AllChannelsCLAHE(clip_limit=(1, 10)),
iaa.CLAHE(clip_limit=(1, 10)),
iaa.Alpha((0.0, 1.0), iaa.AllChannelsHistogramEqualization()),
iaa.Alpha((0.0, 1.0), iaa.HistogramEqualization()),
iaa.pillike.Autocontrast(cutoff=(0, 15.0)),
])),
iaa.Sometimes(0.2,
iaa.OneOf([
iaa.MultiplyBrightness(),
iaa.Add((-40, 40)),
iaa.Multiply((0.5, 1.5)),
iaa.Multiply((0.5, 1.5), per_channel=0.5),
iaa.Sometimes(0.15, iaa.MultiplyAndAddToBrightness(mul=(0.5, 1.5), add=(-30, 30)))
])),
iaa.Sometimes(0.15,
iaa.OneOf([iaa.Dropout(p=(0, 0.1)),
iaa.Dropout2d(p=0.1),
iaa.SaltAndPepper(0.1),
])),
iaa.Sometimes(0.25,
iaa.OneOf([iaa.BlendAlphaElementwise((0, 1.0), iaa.AddToHue(100)),
iaa.BlendAlphaSimplexNoise(iaa.EdgeDetect(1.0), upscale_method="linear"),
iaa.BlendAlphaElementwise([0.25, 0.75], iaa.MedianBlur(5)),
iaa.BlendAlphaSomeColors(iaa.AveragePooling(7), from_colorspace="BGR"),
iaa.BlendAlphaHorizontalLinearGradient(
iaa.TotalDropout(1.0),
min_value=0.2, max_value=0.8),
iaa.BlendAlphaHorizontalLinearGradient(
iaa.AveragePooling(11),
start_at=(0.0, 1.0), end_at=(0.0, 1.0)),
])),
iaa.Sometimes(0.15,
iaa.OneOf([
iaa.Clouds(),
iaa.Fog(),
iaa.Snowflakes(flake_size=(0.1, 0.4), speed=(0.01, 0.05)),
iaa.Rain(speed=(0.1, 0.3))
])),
# iaa.Sometimes(0.15,
# iaa.OneOf([
# iaa.Cartoon(),
# iaa.Cartoon(blur_ksize=3, segmentation_size=1.0,
# saturation=2.0, edge_prevalence=1.0),
# iaa.Superpixels(p_replace=0.5, n_segments=64),
# iaa.Superpixels(p_replace=(0.1, 1.0), n_segments=(16, 128)),
# iaa.UniformVoronoi(250, p_replace=0.9, max_size=None),
# ])),
], random_order=True)
def __call__(self, img):
img = np.array(img)
return self.aug.augment_image(img)
if __name__ == "__main__":
from datasets.data_augment import preproc
loader = FaceDataset(root_path=os.path.join('/home/can/AI_Camera/EfficientFaceNet/data/widerface/train/images'), file_name='label.txt',
preproc=preproc(300, (127, 127, 127)))
print(len(loader))
for i in range(0, len(loader)):
print("\n****")
print(i)
a = loader.__getitem__(i)
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment