Upload app.py

app.py

@@ -0,0 +1,312 @@
+import base64
+import os
+import shutil
+import subprocess
+import uuid
+import zipfile
+
+import cv2
+import numpy as np
+from flask import Flask, jsonify, render_template, request, send_file
+from werkzeug.utils import secure_filename
+
+app = Flask(__name__)
+
+UPLOAD_FOLDER = "/tmp/uploads"
+PROCESSED_FOLDER = "/tmp/processed"
+app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER
+app.config["PROCESSED_FOLDER"] = PROCESSED_FOLDER
+
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+os.makedirs(PROCESSED_FOLDER, exist_ok=True)
+
+# Try imageio_ffmpeg (local dev), fall back to system ffmpeg (HF Spaces / Docker)
+try:
+    import imageio_ffmpeg
+
+    FFMPEG_EXE = imageio_ffmpeg.get_ffmpeg_exe()
+except ImportError:
+    FFMPEG_EXE = shutil.which("ffmpeg") or "ffmpeg"
+
+
+@app.route("/")
+def index():
+    return render_template("index.html")
+
+
+@app.route("/upload", methods=["POST"])
+def upload_files():
+    if "videos" not in request.files:
+        return jsonify({"error": "No files"}), 400
+
+    files = request.files.getlist("videos")
+    saved_files = []
+    frame_url = None
+    orig_w, orig_h = 0, 0
+
+    for idx, file in enumerate(files):
+        if file.filename:
+            filename = secure_filename(file.filename)
+            unique_filename = f"{uuid.uuid4().hex}_{filename}"
+            filepath = os.path.join(app.config["UPLOAD_FOLDER"], unique_filename)
+            file.save(filepath)
+            saved_files.append(unique_filename)
+
+            if idx == 0:
+                cap = cv2.VideoCapture(filepath)
+                orig_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+                orig_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+                ret, frame = cap.read()
+                cap.release()
+
+                if ret:
+                    frame_filename = f"{unique_filename}.jpg"
+                    frame_path = os.path.join(
+                        app.config["UPLOAD_FOLDER"], frame_filename
+                    )
+                    cv2.imwrite(frame_path, frame)
+                    frame_url = f"/uploads/{frame_filename}"
+
+    return jsonify(
+        {
+            "filenames": saved_files,
+            "frame_url": frame_url,
+            "orig_w": orig_w,
+            "orig_h": orig_h,
+        }
+    )
+
+
+@app.route("/uploads/<filename>")
+def serve_upload(filename):
+    return send_file(os.path.join(app.config["UPLOAD_FOLDER"], filename))
+
+
+@app.route("/process", methods=["POST"])
+def process_videos():
+    data = request.json
+    filenames = data.get("filenames", [])
+    method = data.get("method", "blur_heavy")
+    tool = data.get("tool", "box")
+    upscale = data.get("upscale", "none")
+
+    orig_w, orig_h = int(data["orig_w"]), int(data["orig_h"])
+    processed_files = []
+
+    # Handle coordinates based on the tool used
+    if tool == "box":
+        px, py, pw, ph = (
+            float(data["px"]),
+            float(data["py"]),
+            float(data["pw"]),
+            float(data["ph"]),
+        )
+        x = max(0, min(int(px * orig_w), orig_w - 2))
+        y = max(0, min(int(py * orig_h), orig_h - 2))
+        w = max(1, min(int(pw * orig_w), orig_w - x))
+        h = max(1, min(int(ph * orig_h), orig_h - y))
+        mask_path = None
+
+    elif tool == "brush":
+        # Decode the painted canvas into an image
+        mask_data = data.get("mask_b64", "").split(",")[1]
+        img_bytes = base64.b64decode(mask_data)
+        nparr = np.frombuffer(img_bytes, np.uint8)
+        mask_img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)  # Reads RGBA
+
+        # Extract the alpha channel (where the user painted)
+        alpha_channel = mask_img[:, :, 3]
+        full_mask = cv2.resize(
+            alpha_channel, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST
+        )
+
+        # Find the exact bounding box of their painting
+        coords = cv2.findNonZero(full_mask)
+        if coords is None:
+            return jsonify({"error": "No brush strokes detected."}), 400
+
+        x, y, w, h = cv2.boundingRect(coords)
+
+        # Crop the mask to the exact bounding box to map it to FFmpeg
+        mask_crop = full_mask[y : y + h, x : x + w]
+        mask_path = os.path.join(
+            app.config["UPLOAD_FOLDER"], f"mask_{uuid.uuid4().hex}.png"
+        )
+        cv2.imwrite(mask_path, mask_crop)
+
+    for filename in filenames:
+        input_path = os.path.join(app.config["UPLOAD_FOLDER"], filename)
+        # Extract original name by removing the uuid prefix (32 hex chars + underscore)
+        original_name = filename.split("_", 1)[1] if "_" in filename else filename
+        name_part, ext_part = os.path.splitext(original_name)
+        friendly_name = f"{name_part}_Cle{ext_part}"
+        output_filename = f"clean_{filename}"
+        output_path = os.path.join(app.config["PROCESSED_FOLDER"], output_filename)
+
+        # Build FFmpeg commands depending on the tool and method
+        if tool == "box":
+            if method == "delogo":
+                vf = f"delogo=x={x}:y={y}:w={w}:h={h}"
+            elif method == "blur_light":
+                vf = f"split[m][r];[r]crop={w}:{h}:{x}:{y},gblur=sigma=10[b];[m][b]overlay={x}:{y}"
+            elif method == "blur_heavy":
+                vf = f"split[m][r];[r]crop={w}:{h}:{x}:{y},gblur=sigma=35[b];[m][b]overlay={x}:{y}"
+            elif method == "pixelate":
+                pw_sq, ph_sq = max(1, w // 15), max(1, h // 15)
+                vf = f"split[m][r];[r]crop={w}:{h}:{x}:{y},scale={pw_sq}:{ph_sq},scale={w}:{h}:flags=neighbor[p];[m][p]overlay={x}:{y}"
+            elif method == "black_box":
+                vf = f"drawbox=x={x}:y={y}:w={w}:h={h}:color=black:t=fill"
+
+            command = [
+                FFMPEG_EXE,
+                "-y",
+                "-i",
+                input_path,
+                "-vf",
+                vf,
+                "-c:a",
+                "copy",
+                output_path,
+            ]
+
+        elif tool == "brush":
+            # Brush uses an advanced alpha-merge complex filter to shape the blur to the paint strokes
+            if method == "delogo":
+                # Delogo strictly requires a box, so we fall back to the bounding box of the paint
+                command = [
+                    FFMPEG_EXE,
+                    "-y",
+                    "-i",
+                    input_path,
+                    "-vf",
+                    f"delogo=x={x}:y={y}:w={w}:h={h}",
+                    "-c:a",
+                    "copy",
+                    output_path,
+                ]
+            else:
+                if method == "blur_light":
+                    effect = f"[0:v]crop={w}:{h}:{x}:{y},gblur=sigma=10[fx];"
+                elif method == "blur_heavy":
+                    effect = f"[0:v]crop={w}:{h}:{x}:{y},gblur=sigma=35[fx];"
+                elif method == "pixelate":
+                    pw_sq, ph_sq = max(1, w // 15), max(1, h // 15)
+                    effect = f"[0:v]crop={w}:{h}:{x}:{y},scale={pw_sq}:{ph_sq},scale={w}:{h}:flags=neighbor[fx];"
+                elif method == "black_box":
+                    effect = f"color=black:size={w}x{h}[fx];"
+
+                # Combine the cropped effect with the painted alpha mask
+                filter_chain = (
+                    effect
+                    + f"[fx][1:v]alphamerge[al];[0:v][al]overlay={x}:{y}:shortest=1"
+                )
+
+                command = [
+                    FFMPEG_EXE,
+                    "-y",
+                    "-i",
+                    input_path,
+                    "-loop",
+                    "1",
+                    "-i",
+                    mask_path,  # Load the painted mask as an infinite video layer
+                    "-filter_complex",
+                    filter_chain,
+                    "-c:a",
+                    "copy",
+                    output_path,
+                ]
+
+        try:
+            subprocess.run(
+                command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+            )
+            processed_files.append((output_filename, friendly_name))
+        except subprocess.CalledProcessError as e:
+            print("FFMPEG ERROR:", e.stderr.decode())
+            return jsonify({"error": "FFmpeg processing failed"}), 500
+
+    # --- Optional Upscale Pass ---
+    if upscale != "none":
+        upscaled_files = []
+        for internal_name, friendly in processed_files:
+            src = os.path.join(app.config["PROCESSED_FOLDER"], internal_name)
+            up_filename = f"up_{internal_name}"
+            up_path = os.path.join(app.config["PROCESSED_FOLDER"], up_filename)
+
+            # Determine target resolution
+            if upscale == "1.5x":
+                scale_expr = "scale=iw*1.5:ih*1.5:flags=lanczos"
+            elif upscale == "2x":
+                scale_expr = "scale=iw*2:ih*2:flags=lanczos"
+            elif upscale == "4k":
+                scale_expr = "scale=3840:2160:force_original_aspect_ratio=decrease:flags=lanczos,pad=3840:2160:(ow-iw)/2:(oh-ih)/2"
+            else:
+                scale_expr = "scale=iw*2:ih*2:flags=lanczos"
+
+            # Lanczos upscale + light unsharp mask to sharpen
+            vf_upscale = f"{scale_expr},unsharp=5:5:0.5:5:5:0.0"
+
+            up_cmd = [
+                FFMPEG_EXE,
+                "-y",
+                "-i",
+                src,
+                "-vf",
+                vf_upscale,
+                "-c:v",
+                "libx264",
+                "-crf",
+                "18",
+                "-preset",
+                "slow",
+                "-c:a",
+                "copy",
+                up_path,
+            ]
+
+            try:
+                subprocess.run(
+                    up_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+                )
+                # Update friendly name to indicate upscale
+                fn_part, fn_ext = os.path.splitext(friendly)
+                upscaled_friendly = f"{fn_part}_HD{fn_ext}"
+                upscaled_files.append((up_filename, upscaled_friendly))
+            except subprocess.CalledProcessError as e:
+                print("UPSCALE ERROR:", e.stderr.decode())
+                return jsonify({"error": "Video upscaling failed"}), 500
+
+        processed_files = upscaled_files
+
+    if len(processed_files) == 1:
+        internal_name, friendly = processed_files[0]
+        return jsonify(
+            {"download_url": f"/download/{internal_name}", "download_name": friendly}
+        )
+    else:
+        zip_filename = f"bulk_processed_{uuid.uuid4().hex[:6]}.zip"
+        zip_path = os.path.join(app.config["PROCESSED_FOLDER"], zip_filename)
+        with zipfile.ZipFile(zip_path, "w") as zipf:
+            for internal_name, friendly in processed_files:
+                zipf.write(
+                    os.path.join(app.config["PROCESSED_FOLDER"], internal_name),
+                    friendly,
+                )
+        return jsonify(
+            {"download_url": f"/download/{zip_filename}", "download_name": zip_filename}
+        )
+
+
+@app.route("/download/<filename>")
+def download_file(filename):
+    download_name = request.args.get("name", filename)
+    return send_file(
+        os.path.join(app.config["PROCESSED_FOLDER"], filename),
+        as_attachment=True,
+        download_name=download_name,
+    )
+
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=7860, debug=False)