generated from fastai/nbdev_template
-
Notifications
You must be signed in to change notification settings - Fork 29
Expand file tree
/
Copy pathautoencoder.py
More file actions
64 lines (56 loc) · 2.19 KB
/
autoencoder.py
File metadata and controls
64 lines (56 loc) · 2.19 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
# AUTOGENERATED! DO NOT EDIT! File to edit: 07_autoencoder.ipynb (unless otherwise specified).
__all__ = ['get_pixel', 'change_image_background', 'create_augmentor_pipeline', 'load_data']
# Cell
import Augmentor
import os
import numpy as np
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
from PIL import Image
from tqdm import tqdm
# Cell
def get_pixel(image, i, j):
""" Returns a pixel at coordinate (`i`, `j`). """
return image.getpixel((i,j))
def change_image_background(orig_dir_path, converted_path):
""" Changes the image background from white to black and foreground from black to white,
for all the images at folder `orig_dir_path` and place them into folder `converted_path`."""
files = os.listdir(dir_path)
num_files = len(files)
data = []
counter = 1
for f in tqdm(files, total=num_files):
img = Image.open(os.path.join(dir_path,f))
out_img = Image.new('RGB',img.size,color=1)
width, height = img.size
for w in range(width):
for h in range(height):
r, g, b = get_pixel(img, w,h)
if r > 128 or g > 128 or b > 128:
r = g = b = 0
else:
r = g = b = 255
out_img.putpixel((w,h),(r,g,b))
file_name = os.path.join(converted_path, str(counter) + '.png')
out_img.save(file_name)
counter += 1
return data
def create_augmentor_pipeline(dir_path):
""" Creates a pipeline for generating extra images from images at folder `dir_path`."""
p = Augmentor.Pipeline(dir_path)
p.resize(probability=1,width=64,height=64)
p.rotate90(probability=0.1)
p.rotate(probability=0.2, max_left_rotation=5, max_right_rotation=10)
p.skew_left_right(probability=0.1)
p.greyscale(probability=1)
return p
def load_data(dir_path):
""" Loads all the images from directory `dir_path`, converts them to matrices and return a list."""
files = os.listdir(dir_path)
num_files = len(files)
data = []
for f in tqdm(files, total=num_files):
img = Image.open(os.path.join(dir_path,f))
img_array = np.array(img)
data.append(img_array)
return data