import os from datasets import load_dataset import json import uuid from pathlib import Path import json from datasets import load_dataset from flask import Flask from flask_cors import CORS from PIL import Image app = Flask(__name__, static_url_path='/static') CORS(app) TOKEN = os.environ.get('dataset_token') dataset = load_dataset( "huggingface-projects/wordalle_prompts", use_auth_token=TOKEN) Path("static/images").mkdir(parents=True, exist_ok=True) # extract images and prompts from dataset and save to dis data = {} for row in dataset['train']: prompt = dataset['train'].features['label'].int2str(row['label']) image = row['image'] hash = uuid.uuid4().hex image_file = Path(f'static/images/{hash}.jpg') image_compress = image.resize((136, 136), Image.Resampling.LANCZOS) image_compress.save(image_file, optimize=True, quality=95) if prompt not in data: data[prompt] = [] data[prompt].append(str(image_file)) with open('static/data.json', 'w') as f: json.dump(data, f) @app.route('/') def index(): return app.send_static_file('index.html') @app.route('/data') def getdata(): return app.send_static_file('data.json') if __name__ == '__main__': app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 7860)))