mirror of
https://github.com/igv/FSRCNN-TensorFlow.git
synced 2026-08-19 01:22:18 +08:00
73 lines
1.9 KiB
Python
73 lines
1.9 KiB
Python
import os
|
|
import sys
|
|
import glob
|
|
import numpy as np
|
|
from PIL import Image
|
|
import pdb
|
|
from random import randint
|
|
|
|
# Artifically expands the dataset by a factor of 19 by scaling and then rotating every image
|
|
def main():
|
|
if len(sys.argv) == 2:
|
|
data = prepare_data(sys.argv[1])
|
|
else:
|
|
print("Missing argument: You must specify a folder with images to expand")
|
|
return
|
|
|
|
for i in range(len(data)):
|
|
scale(data[i], randint(0, 3))
|
|
rotate(data[i], randint(0, 2))
|
|
c = randint(0, 1)
|
|
if c==0:
|
|
flip_left_right(data[i])
|
|
else:
|
|
flip_top_bottom(data[i])
|
|
|
|
def prepare_data(dataset):
|
|
filenames = os.listdir(dataset)
|
|
data_dir = os.path.join(os.getcwd(), dataset)
|
|
data = glob.glob(os.path.join(data_dir, "*.bmp"))
|
|
|
|
return data
|
|
|
|
def scale(file, scale):
|
|
image = Image.open(file)
|
|
width, height = image.size
|
|
|
|
scales = [0.9, 0.8, 0.7, 0.6]
|
|
new_width, new_height = int(width * scales[scale]), int(height * scales[scale])
|
|
new_image = image.resize((new_width, new_height), Image.BICUBIC)
|
|
new_path = '{}-{}.bmp'.format(file[:-4], scales[scale])
|
|
new_image.save(new_path)
|
|
"""
|
|
rotate(new_path, randint(0, 2))
|
|
c = randint(0, 1)
|
|
if c==0:
|
|
flip_left_right(new_path)
|
|
else:
|
|
flip_top_bottom(new_path)
|
|
"""
|
|
|
|
def rotate(file, rotation):
|
|
image = Image.open(file)
|
|
|
|
rotations = [90, 180, 270]
|
|
new_image = image.rotate(rotations[rotation], expand=True)
|
|
new_path = '{}-{}.bmp'.format(file[:-4], rotations[rotation])
|
|
new_image.save(new_path)
|
|
|
|
def flip_left_right(file):
|
|
image = Image.open(file)
|
|
new_image = image.transpose(Image.FLIP_LEFT_RIGHT)
|
|
new_path = '{}-{}.bmp'.format(file[:-4], "lr")
|
|
new_image.save(new_path)
|
|
|
|
def flip_top_bottom(file):
|
|
image = Image.open(file)
|
|
new_image = image.transpose(Image.FLIP_TOP_BOTTOM)
|
|
new_path = '{}-{}.bmp'.format(file[:-4], "tb")
|
|
new_image.save(new_path)
|
|
|
|
if __name__ == '__main__':
|
|
main()
|