import streamlit as st
from PIL import Image
from os.path import join
from os import listdir
from src.trainer import GANTrainer
from src.miscc.config import cfg
from src.miscc.config import cfg, cfg_from_file
MODEL_FILE = '/home/sondn/DIY/StackGAN-Pytorch/output/res/step2.pth'
IMGS_DIR = '/home/sondn/DIY/StackGAN-Pytorch/output/images'
cfg_from_file('src/cfg/vn-celeb_eval.yml')
st.title('Text-to-face app')
desc = st.text_area('Type your face description here')
ok = st.button('OK')
if ok:
algo = GANTrainer('')
algo.sample_one(desc=desc, samples=25, save_dir=IMGS_DIR, truncate_dir=True)
columns = st.beta_columns(5)
next = st.button('Next')
if next:
for idx, img_file in enumerate(listdir(IMGS_DIR)):
with columns[idx % 5]:
img = Image.open(join(IMGS_DIR, img_file))
st.image(img, channels='BGR', use_column_width=True)
start_epoch = int(start_epoch) + 1
print('start_epoch', start_epoch)
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
discriminator = discriminator.cuda()
labels = labels.cuda()
return text_encoder, image_encoder, discriminator, labels, start_epoch
if __name__ == "__main__":
args = parse_args()
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.gpu_id == -1:
cfg.CUDA = False
else:
cfg.GPU_ID = args.gpu_id
if args.data_dir != '':
cfg.DATA_DIR = args.data_dir
print('Using config:')
pprint.pprint(cfg)
if not cfg.TRAIN.FLAG:
args.manualSeed = 100
elif args.manualSeed is None:
args.manualSeed = random.randint(1, 10000)
with open(embedding_file, 'wb') as f:
pickle.dump(features, f)
f.close()
return features
def embed_all_text_data(self, embedding_file=None):
text_dir = join(cfg.DATA_DIR, 'text')
descs = []
for sub_dir in listdir(text_dir):
for txt in listdir(join(text_dir, sub_dir)):
txt_file_path = join(text_dir, sub_dir, txt)
lines = []
with open(txt_file_path, 'r') as f:
for line in f:
lines.append(line)
desc = ViTokenizer.tokenize('.'.join(lines))
descs.append((int(txt.split('.')[0]), desc))
descs.sort(key=lambda tup: tup[0])
sorted_desc = map(lambda tup: tup[1], descs)
embedding_file = join(
text_dir,
'embedding.pkl') if embedding_file is None else embedding_file
self.embed_text(sorted_desc, embedding_file=embedding_file)
if __name__ == '__main__':
cfg_from_file('/home/sondn/DIY/StackGAN-Pytorch/src/cfg/vn-celeb_eval.yml')
embed = TextEmbedding()
embed.embed_all_text_data('/home/sondn/DIY/StackGAN-Pytorch/data/emb.pkl')