mirror of
https://github.com/prise6/smart-iss-posts
synced 2024-04-27 11:31:51 +02:00
80 lines
2.4 KiB
Python
80 lines
2.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
import os
|
|
import random
|
|
import shutil
|
|
import numpy as np
|
|
import re
|
|
|
|
class CollectionManagerFromDirectory:
|
|
|
|
def __init__(self, config, sampling_type = 'autoencoder'):
|
|
self.config = config
|
|
config_sampling = self.config.get('sampling')[sampling_type]
|
|
self.dir = self.config.get('directory')[config_sampling['directory']['from']]
|
|
|
|
jpg_regex = re.compile(".*jpg$")
|
|
self.pictures_id = [pict for pict in os.listdir(self.dir) if jpg_regex.match(pict)]
|
|
|
|
self.dir_base = config_sampling['directory']['base']
|
|
self.dir_train = config_sampling['directory']['train']
|
|
self.dir_test = config_sampling['directory']['test']
|
|
self.dir_valid = config_sampling['directory']['valid']
|
|
|
|
self.seed = config_sampling['seed']
|
|
self.proportions = config_sampling['proportions']
|
|
self.volumes = {}
|
|
|
|
self.shuffle()
|
|
|
|
def count(self):
|
|
self.volumes['total'] = len(self.pictures_id)
|
|
return self
|
|
|
|
def shuffle(self):
|
|
random.seed(self.seed)
|
|
random.shuffle(self.pictures_id)
|
|
return self
|
|
|
|
@staticmethod
|
|
def create_dir(path):
|
|
if os.path.exists(path):
|
|
shutil.rmtree(path, ignore_errors=True)
|
|
|
|
return os.makedirs(path)
|
|
|
|
@staticmethod
|
|
def copy_pictures(dest, picture_dir, pictures_id):
|
|
|
|
pictures_src = os.path.join(picture_dir, pictures_id)
|
|
pictures_dest = os.path.join(dest, pictures_id)
|
|
return shutil.copyfile(pictures_src, pictures_dest)
|
|
|
|
def createDirectories(self):
|
|
CollectionManagerFromDirectory.create_dir(self.dir_base)
|
|
CollectionManagerFromDirectory.create_dir(self.dir_train)
|
|
CollectionManagerFromDirectory.create_dir(self.dir_test)
|
|
CollectionManagerFromDirectory.create_dir(self.dir_valid)
|
|
|
|
return self
|
|
|
|
|
|
def populateDirectories(self):
|
|
self.volumes['train'] = int(np.floor(self.proportions['train'] * self.volumes['total']))
|
|
self.volumes['test'] = int(np.floor(self.proportions['test'] * self.volumes['total']))
|
|
self.volumes['valid'] = self.volumes['total'] - (self.volumes['train'] + self.volumes['test'])
|
|
|
|
for pict in self.pictures_id[:self.volumes['train']]:
|
|
CollectionManagerFromDirectory.copy_pictures(self.dir_train, self.dir, pict)
|
|
|
|
for pict in self.pictures_id[self.volumes['train']:(self.volumes['train'] + self.volumes['test'])]:
|
|
CollectionManagerFromDirectory.copy_pictures(self.dir_test, self.dir, pict)
|
|
|
|
for pict in self.pictures_id[-self.volumes['valid']:]:
|
|
CollectionManagerFromDirectory.copy_pictures(self.dir_valid, self.dir, pict)
|
|
|
|
|
|
|
|
|
|
|