import os, uuid, zipfile, shutil, cv2, json
from flask import Flask, request, send_file, render_template_string, jsonify
from rembg import remove, new_session
from PIL import Image
app = Flask(__name__)
session = new_session("u2net")
JOBS_DIR = "jobs"
os.makedirs(JOBS_DIR, exist_ok=True)
HTML = """
Xlnk Video Processor
"""
@app.route("/")
def index(): return render_template_string(HTML)
@app.route("/upload", methods=["POST"])
def upload():
file = request.files['video']
job_id = str(uuid.uuid4())
path = os.path.join(JOBS_DIR, job_id)
os.makedirs(os.path.join(path, "frames"), exist_ok=True)
video_path = os.path.join(path, "input.mp4")
file.save(video_path)
# Save initial status
with open(os.path.join(path, "status.json"), "w") as f:
json.dump({"status": "processing", "current": 0, "total": 0, "percent": 0}, f)
# In a real app, we'd use a background thread here.
# For a simple HF Space, we'll process it in the route but
# update the file so the frontend can read it.
process_video(job_id, video_path, path)
return jsonify({"job_id": job_id})
def process_video(job_id, video_path, workspace):
cap = cv2.VideoCapture(video_path)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
frames_dir = os.path.join(workspace, "frames")
count = 0
while cap.isOpened():
ret, frame = cap.read()
if not ret: break
output = remove(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), session=session)
Image.fromarray(output).save(os.path.join(frames_dir, f"f_{count:04d}.png"))
count += 1
# Update status file every 5 frames to save disk I/O
if count % 5 == 0 or count == total_frames:
with open(os.path.join(workspace, "status.json"), "w") as f:
json.dump({"status": "processing", "current": count, "total": total_frames, "percent": int((count/total_frames)*100)}, f)
cap.release()
# Zip it up
zip_path = os.path.join(JOBS_DIR, f"{job_id}.zip")
with zipfile.ZipFile(zip_path, 'w') as z:
for f in os.listdir(frames_dir):
z.write(os.path.join(frames_dir, f), f)
with open(os.path.join(workspace, "status.json"), "w") as f:
json.dump({"status": "completed"}, f)
@app.route("/status/")
def status(job_id):
try:
with open(os.path.join(JOBS_DIR, job_id, "status.json"), "r") as f:
return jsonify(json.load(f))
except: return jsonify({"status": "not_found"})
@app.route("/download/")
def download(job_id):
return send_file(os.path.join(JOBS_DIR, f"{job_id}.zip"), as_attachment=True)
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860)