90 lines
3.2 KiB
Python
90 lines
3.2 KiB
Python
import cv2
|
|
import os
|
|
import random
|
|
from PIL import Image
|
|
|
|
def rename_files(directory_path: str):
|
|
extensions = set()
|
|
|
|
for file in os.listdir(directory_path):
|
|
filename, extension = file.split(".")
|
|
extensions.add(extension)
|
|
old_filename = os.path.join(directory_path, file)
|
|
new_filename = os.path.join(directory_path, f"{filename}.jpg")
|
|
if old_filename != new_filename:
|
|
print(f"{old_filename} -> {new_filename}")
|
|
os.rename(old_filename, new_filename)
|
|
|
|
def check_images(directory_path: str):
|
|
for file in os.listdir(directory_path):
|
|
try:
|
|
img = cv2.imread(os.path.join(directory_path, file))
|
|
if img is None:
|
|
print(file)
|
|
except Exception:
|
|
print(file)
|
|
|
|
def check_labels(dir_path: str):
|
|
for filename in os.listdir(dir_path):
|
|
name, extension = filename.split(".")
|
|
if len(name) != 17:
|
|
print(filename)
|
|
|
|
def check_symbols(dir_path: str):
|
|
with open(os.path.join(dir_path, "dict.txt"), "r") as dict_file:
|
|
dict_content = dict_file.readlines()
|
|
dict_chars = set(char.strip() for char in dict_content)
|
|
for filename in os.listdir(os.path.join(dir_path, "images")):
|
|
label, extension = filename.split(".")
|
|
if any([char not in dict_chars for char in label]):
|
|
print(filename)
|
|
|
|
def max_height(dir_path: str):
|
|
max_height = 0
|
|
max_filename = ''
|
|
for filename in os.listdir(dir_path):
|
|
im = Image.open(os.path.join(dir_path, filename))
|
|
if im.height > max_height:
|
|
max_height = im.height
|
|
max_filename = filename
|
|
|
|
print(max_filename, max_height)
|
|
|
|
def resize_to_height(dir_path: str, target_height=48):
|
|
for filename in os.listdir(dir_path):
|
|
with Image.open(os.path.join(dir_path, filename)) as img:
|
|
width_percent = target_height / float(img.height)
|
|
new_width = int(float(img.width) * width_percent)
|
|
resized_img = img.resize((new_width, target_height), Image.LANZOS)
|
|
resized_img.save(os.path.join(dir_path, filename))
|
|
|
|
def split_dataset(dir_path: str):
|
|
images_dir = os.path.join(dir_path, "images")
|
|
filenames = os.listdir(images_dir)
|
|
dataset_length = len(filenames)
|
|
random.shuffle(filenames)
|
|
train_ratio = 0.8
|
|
val_ratio = 0.2
|
|
train_files_len = round(dataset_length * train_ratio)
|
|
val_files_len = round(dataset_length * val_ratio)
|
|
train_files = filenames[:train_files_len]
|
|
val_files = filenames[train_files_len:train_files_len+val_files_len]
|
|
train_file_path = os.path.join(dir_path, "train.txt")
|
|
val_file_path = os.path.join(dir_path, "val.txt")
|
|
with open(train_file_path, "w") as train_file:
|
|
for filename in train_files:
|
|
label, _ = filename.split(".")
|
|
train_file.write(f"{filename}\t{label}\n")
|
|
|
|
with open(val_file_path, "w") as val_file:
|
|
for filename in val_files:
|
|
label, _ = filename.split(".")
|
|
val_file.write(f"{filename}\t{label}\n")
|
|
|
|
# rename_files("train_data/images/")
|
|
# check_images("train_data/images/")
|
|
# check_labels("train_data/images/")
|
|
# check_symbols("train_data/")
|
|
# max_height("train_data/images")
|
|
split_dataset("train_data/")
|