diff --git a/Engine/Model/convert_full.py b/Engine/Model/convert_full.py new file mode 100644 index 0000000..4d8121b --- /dev/null +++ b/Engine/Model/convert_full.py @@ -0,0 +1,55 @@ +from PIL import Image +import numpy, pefile, os + +#################################################################################################### + +def preprocess_image(file_data, size, channels=1): + width, height = size + wah = int(numpy.ceil(numpy.sqrt(len(file_data) / channels))) + arr = numpy.frombuffer(file_data, dtype=numpy.uint8) + img = numpy.zeros(wah*wah*channels, dtype=numpy.uint8) + img[:len(file_data)] = arr + image = Image.fromarray(img.reshape((wah, wah)), 'L') + return image.resize((width, height), Image.Resampling.NEAREST) + +#################################################################################################### + +def get_type(file_path): + suffix = {".com", ".dll", ".drv", ".exe", ".ocx", ".scr", ".sys", ".mui", ".cpl"} + try: + file_content = None + if not os.path.splitext(file_path)[-1].lower() in suffix: + return file_content + with open(file_path, 'rb') as f: + file_content = f.read() + except: + pass + return file_content + +#################################################################################################### + +def file_to_images(input_dir, output_dir): + for root, dirs, files in os.walk(input_dir): + for file in files: + try: + file_path = os.path.join(root, file) + file_content = get_type(file_path) + if file_content: + print(file_path) + relative_path = os.path.relpath(root, input_dir) + output_path_dir = os.path.join(output_dir, relative_path) + + if not os.path.exists(output_path_dir): + os.makedirs(output_path_dir) + image = preprocess_image(file_content, (224, 224)) + safe_name = os.path.basename(file_path).replace('/', '_').replace('\\', '_') + image.save(os.path.join(output_path_dir, f"{safe_name}.png")) + except Exception as e: + print(e) + +#################################################################################################### + +input_directory = input('Enter the folder path: ') +output_directory = "./Image_Output" +file_to_images(input_directory, output_directory) +input('Conversion Complete') diff --git a/Engine/Model/train.py b/Engine/Model/train.py index 4592690..09f7e20 100644 --- a/Engine/Model/train.py +++ b/Engine/Model/train.py @@ -11,7 +11,7 @@ tf.keras.mixed_precision.set_global_policy('mixed_float16') #################################################################################################### -def load_dataset(data_dir, image_size=(224, 224), val_split=0.01, batch_size=256, color_mode="grayscale"): +def load_dataset(data_dir, image_size=(224, 224), val_split=0.00001, batch_size=256, color_mode="grayscale"): channels = 1 if color_mode == "grayscale" else 3 file_paths, labels, class_indices = get_file_list(data_dir) @@ -142,7 +142,7 @@ class CustomModelCheckpoint(callbacks.Callback): #################################################################################################### -train_ds, val_ds = load_dataset(r'.\Image_File_Pefile') +train_ds, val_ds = load_dataset(r'.\Image_Full_Pefile') for imgs, _ in train_ds.take(1): train_ds.image_shape = imgs.shape[1:] break