131 lines
4.1 KiB
Python
131 lines
4.1 KiB
Python
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from flask import Flask, render_template, request, redirect, url_for, Response, send_from_directory
|
|
from werkzeug.middleware.proxy_fix import ProxyFix
|
|
|
|
from db import init_db, search_images
|
|
|
|
app = Flask(__name__, static_folder='Pictures', static_url_path='/pictures')
|
|
app.wsgi_app = ProxyFix(app.wsgi_app, x_prefix=1)
|
|
init_db()
|
|
|
|
# job_id -> {'done': int, 'total': int, 'finished': bool, 'site': str, 'tags': str}
|
|
downloads = {}
|
|
|
|
|
|
@app.route('/previews/<path:filename>')
|
|
def serve_preview(filename):
|
|
return send_from_directory('Previews', filename)
|
|
|
|
|
|
@app.route('/')
|
|
def slideshow():
|
|
raw_query = request.args.get('tags', '').strip()
|
|
results = search_images(raw_query)
|
|
base_dir = os.path.dirname(os.path.abspath(__file__))
|
|
pictures_dir = os.path.join(base_dir, 'Pictures')
|
|
previews_dir = os.path.join(base_dir, 'Previews')
|
|
|
|
image_urls = []
|
|
preview_urls = []
|
|
post_urls = []
|
|
tags_list = []
|
|
file_sizes = []
|
|
preview_sizes = []
|
|
|
|
for r in results:
|
|
fname = r['filename']
|
|
pname = r['preview_filename']
|
|
image_urls.append(f'pictures/{fname}')
|
|
post_urls.append(r['post_url'])
|
|
tags_list.append(r['tags'].split())
|
|
file_sizes.append(os.path.getsize(os.path.join(pictures_dir, fname)))
|
|
|
|
ppath = os.path.join(previews_dir, pname) if pname else ''
|
|
if pname and os.path.exists(ppath):
|
|
preview_urls.append(f'previews/{pname}')
|
|
preview_sizes.append(os.path.getsize(ppath))
|
|
else:
|
|
preview_urls.append(f'pictures/{fname}')
|
|
preview_sizes.append(file_sizes[-1])
|
|
|
|
active_tags = raw_query.split() if raw_query else []
|
|
job_id = request.args.get('job_id')
|
|
return render_template(
|
|
'slideshow.html',
|
|
images=image_urls,
|
|
preview_images=preview_urls,
|
|
post_urls=post_urls,
|
|
tags_list=tags_list,
|
|
file_sizes=file_sizes,
|
|
preview_sizes=preview_sizes,
|
|
active_tags=active_tags,
|
|
tag_query=raw_query,
|
|
job_id=job_id,
|
|
)
|
|
|
|
|
|
@app.route('/download', methods=['POST'])
|
|
def download():
|
|
tags = request.form.get('tags', '').strip()
|
|
site = request.form.get('site', 'e621')
|
|
if not tags:
|
|
return redirect(url_for('slideshow', tags=tags))
|
|
|
|
job_id = uuid.uuid4().hex[:8]
|
|
downloads[job_id] = {'done': 0, 'total': 0, 'finished': False, 'site': site, 'tags': tags, 'status': ''}
|
|
|
|
proc = subprocess.Popen(
|
|
[sys.executable, 'downloader.py', '--site', site, '--query', tags],
|
|
cwd=app.root_path,
|
|
stdout=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
|
|
def read_stdout():
|
|
for line in proc.stdout:
|
|
line = line.strip()
|
|
if line.startswith('status:'):
|
|
downloads[job_id]['status'] = line[7:]
|
|
elif line.startswith('total:'):
|
|
downloads[job_id]['total'] = int(line.split(':')[1])
|
|
elif line.startswith('progress:'):
|
|
done, total = line.split(':')[1].split('/')
|
|
downloads[job_id]['done'] = int(done)
|
|
downloads[job_id]['total'] = int(total)
|
|
elif line == 'done':
|
|
downloads[job_id]['finished'] = True
|
|
downloads[job_id]['finished'] = True
|
|
|
|
threading.Thread(target=read_stdout, daemon=True).start()
|
|
return redirect(url_for('slideshow', tags=tags, job_id=job_id))
|
|
|
|
|
|
@app.route('/download/progress/<job_id>')
|
|
def download_progress(job_id):
|
|
def generate():
|
|
while True:
|
|
info = downloads.get(job_id)
|
|
if not info:
|
|
yield f'data: {json.dumps({"error": "not found"})}\n\n'
|
|
break
|
|
yield f'data: {json.dumps({"done": info["done"], "total": info["total"], "finished": info["finished"], "status": info["status"]})}\n\n'
|
|
if info['finished']:
|
|
break
|
|
time.sleep(0.3)
|
|
|
|
return Response(
|
|
generate(),
|
|
mimetype='text/event-stream',
|
|
headers={'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'},
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
app.run(host='0.0.0.0', debug=False)
|