Spaces:
Sleeping
Sleeping
| import os | |
| import cv2 | |
| import json | |
| import shutil | |
| import mimetypes | |
| import tempfile | |
| import subprocess | |
| import numpy as np | |
| from fastapi import ( | |
| FastAPI, | |
| UploadFile, | |
| File, | |
| Form, | |
| HTTPException | |
| ) | |
| from fastapi.responses import ( | |
| HTMLResponse, | |
| FileResponse | |
| ) | |
| from starlette.background import BackgroundTask | |
| # ========================================================= | |
| # APP | |
| # ========================================================= | |
| app = FastAPI( | |
| title="Fast Watermark Remover" | |
| ) | |
| # ========================================================= | |
| # CONFIG | |
| # ========================================================= | |
| MAX_FILE_SIZE = 500 * 1024 * 1024 | |
| PREVIEW_WIDTH = 720 | |
| PREVIEW_HEIGHT = 1280 | |
| PROCESS_WIDTH = 720 | |
| ALLOWED_EXTENSIONS = [ | |
| ".mp4", | |
| ".mov", | |
| ".avi", | |
| ".mkv", | |
| ".webm", | |
| ".flv", | |
| ".wmv", | |
| ".mpeg", | |
| ".mpg", | |
| ".m4v", | |
| ".3gp" | |
| ] | |
| # ========================================================= | |
| # HELPERS | |
| # ========================================================= | |
| def cleanup(path): | |
| try: | |
| if os.path.exists(path): | |
| shutil.rmtree(path) | |
| except: | |
| pass | |
| def allowed_file(filename): | |
| ext = os.path.splitext( | |
| filename | |
| )[1].lower() | |
| return ext in ALLOWED_EXTENSIONS | |
| # ========================================================= | |
| # MASK | |
| # ========================================================= | |
| def create_mask( | |
| width, | |
| height, | |
| brush_points | |
| ): | |
| mask = np.zeros( | |
| (height, width), | |
| dtype=np.uint8 | |
| ) | |
| for point in brush_points: | |
| x = int(point["x"]) | |
| y = int(point["y"]) | |
| size = int(point["size"]) | |
| cv2.circle( | |
| mask, | |
| (x, y), | |
| size, | |
| 255, | |
| -1 | |
| ) | |
| mask = cv2.GaussianBlur( | |
| mask, | |
| (15, 15), | |
| 0 | |
| ) | |
| return mask | |
| def remove_watermark( | |
| frame, | |
| brush_points | |
| ): | |
| h, w = frame.shape[:2] | |
| mask = create_mask( | |
| w, | |
| h, | |
| brush_points | |
| ) | |
| result = cv2.inpaint( | |
| frame, | |
| mask, | |
| 5, | |
| cv2.INPAINT_TELEA | |
| ) | |
| return result | |
| # ========================================================= | |
| # UI | |
| # ========================================================= | |
| async def home(): | |
| return """ | |
| <!DOCTYPE html> | |
| <html> | |
| <head> | |
| <title>Fast Watermark Remover</title> | |
| <meta | |
| name="viewport" | |
| content="width=device-width, initial-scale=1.0" | |
| /> | |
| <style> | |
| *{ | |
| margin:0; | |
| padding:0; | |
| box-sizing:border-box; | |
| } | |
| body{ | |
| background:#f5f5f5; | |
| font-family:Arial; | |
| color:#111827; | |
| } | |
| .container{ | |
| max-width:1500px; | |
| margin:auto; | |
| padding:20px; | |
| } | |
| .title{ | |
| font-size:42px; | |
| font-weight:bold; | |
| margin-bottom:8px; | |
| } | |
| .subtitle{ | |
| color:#6b7280; | |
| margin-bottom:25px; | |
| } | |
| .main{ | |
| display:flex; | |
| gap:20px; | |
| align-items:flex-start; | |
| } | |
| .left{ | |
| flex:1; | |
| background:white; | |
| padding:20px; | |
| border-radius:20px; | |
| } | |
| .right{ | |
| width:340px; | |
| background:white; | |
| padding:20px; | |
| border-radius:20px; | |
| } | |
| .video-grid{ | |
| display:flex; | |
| gap:20px; | |
| justify-content:center; | |
| } | |
| .video-box{ | |
| flex:1; | |
| text-align:center; | |
| } | |
| .video-title{ | |
| font-weight:bold; | |
| margin-bottom:12px; | |
| } | |
| .preview-wrapper{ | |
| position:relative; | |
| width:360px; | |
| height:640px; | |
| margin:auto; | |
| border-radius:18px; | |
| overflow:hidden; | |
| background:black; | |
| } | |
| video{ | |
| width:100%; | |
| height:100%; | |
| object-fit:contain; | |
| background:black; | |
| } | |
| canvas{ | |
| position:absolute; | |
| top:0; | |
| left:0; | |
| width:100%; | |
| height:100%; | |
| cursor:crosshair; | |
| } | |
| button{ | |
| width:100%; | |
| border:none; | |
| padding:16px; | |
| border-radius:12px; | |
| cursor:pointer; | |
| font-size:16px; | |
| margin-top:14px; | |
| font-weight:bold; | |
| } | |
| .upload-btn{ | |
| background:#111827; | |
| color:white; | |
| } | |
| .brush-btn{ | |
| background:#10b981; | |
| color:white; | |
| } | |
| .remove-btn{ | |
| background:#4361ee; | |
| color:white; | |
| } | |
| .clear-btn{ | |
| background:#ef4444; | |
| color:white; | |
| } | |
| .download-btn{ | |
| background:#111827; | |
| color:white; | |
| display:none; | |
| } | |
| .slider{ | |
| width:100%; | |
| margin-top:20px; | |
| } | |
| .status{ | |
| margin-top:18px; | |
| font-size:15px; | |
| color:#374151; | |
| } | |
| .loader{ | |
| display:none; | |
| margin-top:20px; | |
| text-align:center; | |
| } | |
| .spinner{ | |
| width:45px; | |
| height:45px; | |
| border:5px solid #ddd; | |
| border-top:5px solid #111827; | |
| border-radius:50%; | |
| margin:auto; | |
| animation:spin 1s linear infinite; | |
| } | |
| @keyframes spin{ | |
| 100%{ | |
| transform:rotate(360deg); | |
| } | |
| } | |
| @media(max-width:1100px){ | |
| .main{ | |
| flex-direction:column; | |
| } | |
| .right{ | |
| width:100%; | |
| } | |
| .video-grid{ | |
| flex-direction:column; | |
| } | |
| .preview-wrapper{ | |
| width:100%; | |
| height:520px; | |
| } | |
| } | |
| </style> | |
| </head> | |
| <body> | |
| <div class="container"> | |
| <div class="title"> | |
| Fast Watermark Remover | |
| </div> | |
| <div class="subtitle"> | |
| Brush paint watermark area and remove it instantly | |
| </div> | |
| <div class="main"> | |
| <div class="left"> | |
| <div class="video-grid"> | |
| <div class="video-box"> | |
| <div class="video-title"> | |
| Original | |
| </div> | |
| <div class="preview-wrapper"> | |
| <video | |
| id="originalVideo" | |
| controls | |
| ></video> | |
| <canvas id="canvas"></canvas> | |
| </div> | |
| </div> | |
| <div class="video-box"> | |
| <div class="video-title"> | |
| Processed | |
| </div> | |
| <div class="preview-wrapper"> | |
| <video | |
| id="outputVideo" | |
| controls | |
| ></video> | |
| </div> | |
| </div> | |
| </div> | |
| </div> | |
| <div class="right"> | |
| <input | |
| type="file" | |
| id="videoInput" | |
| accept="video/*" | |
| hidden | |
| > | |
| <button | |
| class="upload-btn" | |
| onclick="document.getElementById('videoInput').click()" | |
| > | |
| Upload Video | |
| </button> | |
| <button | |
| class="brush-btn" | |
| id="brushToggle" | |
| > | |
| Brush Enabled | |
| </button> | |
| <input | |
| type="range" | |
| min="5" | |
| max="80" | |
| value="25" | |
| class="slider" | |
| id="brushSize" | |
| > | |
| <button | |
| class="remove-btn" | |
| id="removeBtn" | |
| > | |
| Remove Watermark | |
| </button> | |
| <button | |
| class="clear-btn" | |
| id="clearBtn" | |
| > | |
| Clear Brush | |
| </button> | |
| <div | |
| class="status" | |
| id="status" | |
| > | |
| Upload video and paint watermark | |
| </div> | |
| <div | |
| class="loader" | |
| id="loader" | |
| > | |
| <div class="spinner"></div> | |
| <p style="margin-top:15px;"> | |
| Processing video... | |
| </p> | |
| </div> | |
| <a id="downloadLink"> | |
| <button | |
| class="download-btn" | |
| id="downloadBtn" | |
| > | |
| Download Video | |
| </button> | |
| </a> | |
| </div> | |
| </div> | |
| </div> | |
| <script> | |
| const canvas = | |
| document.getElementById( | |
| "canvas" | |
| ); | |
| const ctx = | |
| canvas.getContext("2d"); | |
| const originalVideo = | |
| document.getElementById( | |
| "originalVideo" | |
| ); | |
| const outputVideo = | |
| document.getElementById( | |
| "outputVideo" | |
| ); | |
| const videoInput = | |
| document.getElementById( | |
| "videoInput" | |
| ); | |
| const loader = | |
| document.getElementById( | |
| "loader" | |
| ); | |
| const statusText = | |
| document.getElementById( | |
| "status" | |
| ); | |
| const brushSizeSlider = | |
| document.getElementById( | |
| "brushSize" | |
| ); | |
| let drawing = false; | |
| let brushEnabled = true; | |
| let brushPoints = []; | |
| function resizeCanvas(){ | |
| canvas.width = | |
| canvas.offsetWidth; | |
| canvas.height = | |
| canvas.offsetHeight; | |
| } | |
| window.addEventListener( | |
| "resize", | |
| resizeCanvas | |
| ); | |
| setTimeout( | |
| resizeCanvas, | |
| 500 | |
| ); | |
| document.getElementById( | |
| "brushToggle" | |
| ).onclick = () => { | |
| brushEnabled = | |
| !brushEnabled; | |
| document.getElementById( | |
| "brushToggle" | |
| ).innerText = | |
| brushEnabled | |
| ? "Brush Enabled" | |
| : "Brush Disabled"; | |
| }; | |
| canvas.addEventListener( | |
| "mousedown", | |
| () => { | |
| if(!brushEnabled){ | |
| return; | |
| } | |
| drawing = true; | |
| } | |
| ); | |
| canvas.addEventListener( | |
| "mouseup", | |
| () => { | |
| drawing = false; | |
| ctx.beginPath(); | |
| } | |
| ); | |
| canvas.addEventListener( | |
| "mouseleave", | |
| () => { | |
| drawing = false; | |
| ctx.beginPath(); | |
| } | |
| ); | |
| canvas.addEventListener( | |
| "mousemove", | |
| draw | |
| ); | |
| function draw(e){ | |
| if(!drawing || !brushEnabled){ | |
| return; | |
| } | |
| const rect = | |
| canvas.getBoundingClientRect(); | |
| const x = | |
| e.clientX - rect.left; | |
| const y = | |
| e.clientY - rect.top; | |
| const size = | |
| parseInt( | |
| brushSizeSlider.value | |
| ); | |
| brushPoints.push({ | |
| x:x, | |
| y:y, | |
| size:size, | |
| canvasWidth:canvas.width, | |
| canvasHeight:canvas.height | |
| }); | |
| ctx.lineWidth = size; | |
| ctx.lineCap = "round"; | |
| ctx.strokeStyle = | |
| "rgba(255,0,0,1)"; | |
| ctx.lineTo(x, y); | |
| ctx.stroke(); | |
| ctx.beginPath(); | |
| ctx.moveTo(x, y); | |
| } | |
| document.getElementById( | |
| "clearBtn" | |
| ).onclick = () => { | |
| ctx.clearRect( | |
| 0, | |
| 0, | |
| canvas.width, | |
| canvas.height | |
| ); | |
| brushPoints = []; | |
| }; | |
| videoInput.addEventListener( | |
| "change", | |
| e => { | |
| const file = | |
| e.target.files[0]; | |
| if(!file){ | |
| return; | |
| } | |
| const url = | |
| URL.createObjectURL(file); | |
| originalVideo.src = url; | |
| setTimeout( | |
| resizeCanvas, | |
| 1000 | |
| ); | |
| } | |
| ); | |
| document.getElementById( | |
| "removeBtn" | |
| ).onclick = async () => { | |
| const file = | |
| videoInput.files[0]; | |
| if(!file){ | |
| alert( | |
| "Upload video first" | |
| ); | |
| return; | |
| } | |
| if(brushPoints.length === 0){ | |
| alert( | |
| "Paint watermark first" | |
| ); | |
| return; | |
| } | |
| loader.style.display = | |
| "block"; | |
| statusText.innerText = | |
| "Processing..."; | |
| const formData = | |
| new FormData(); | |
| formData.append( | |
| "file", | |
| file | |
| ); | |
| formData.append( | |
| "brush_points", | |
| JSON.stringify( | |
| brushPoints | |
| )); | |
| try{ | |
| const response = | |
| await fetch( | |
| "/remove/", | |
| { | |
| method:"POST", | |
| body:formData | |
| } | |
| ); | |
| if(!response.ok){ | |
| throw new Error( | |
| "Processing failed" | |
| ); | |
| } | |
| const blob = | |
| await response.blob(); | |
| const outputURL = | |
| URL.createObjectURL(blob); | |
| outputVideo.src = | |
| outputURL; | |
| document.getElementById( | |
| "downloadLink" | |
| ).href = | |
| outputURL; | |
| document.getElementById( | |
| "downloadBtn" | |
| ).style.display = | |
| "block"; | |
| statusText.innerText = | |
| "Done"; | |
| } | |
| catch(error){ | |
| console.error(error); | |
| statusText.innerText = | |
| "Error processing video"; | |
| } | |
| loader.style.display = | |
| "none"; | |
| }; | |
| </script> | |
| </body> | |
| </html> | |
| """ | |
| # ========================================================= | |
| # REMOVE API | |
| # ========================================================= | |
| async def remove_video( | |
| file: UploadFile = File(...), | |
| brush_points: str = Form(...) | |
| ): | |
| if not allowed_file( | |
| file.filename | |
| ): | |
| raise HTTPException( | |
| status_code=400, | |
| detail="Unsupported file" | |
| ) | |
| brush_points = json.loads( | |
| brush_points | |
| ) | |
| session_dir = tempfile.mkdtemp() | |
| original_ext = os.path.splitext( | |
| file.filename | |
| )[1] | |
| original_name = os.path.splitext( | |
| file.filename | |
| )[0] | |
| output_filename = ( | |
| original_name + | |
| "_cleaned" + | |
| original_ext | |
| ) | |
| input_path = os.path.join( | |
| session_dir, | |
| file.filename | |
| ) | |
| output_path = os.path.join( | |
| session_dir, | |
| output_filename | |
| ) | |
| size = 0 | |
| with open(input_path, "wb") as f: | |
| while True: | |
| chunk = await file.read( | |
| 1024 * 1024 | |
| ) | |
| if not chunk: | |
| break | |
| size += len(chunk) | |
| if size > MAX_FILE_SIZE: | |
| cleanup(session_dir) | |
| raise HTTPException( | |
| status_code=400, | |
| detail="File too large" | |
| ) | |
| f.write(chunk) | |
| cap = cv2.VideoCapture( | |
| input_path | |
| ) | |
| width = int( | |
| cap.get( | |
| cv2.CAP_PROP_FRAME_WIDTH | |
| ) | |
| ) | |
| height = int( | |
| cap.get( | |
| cv2.CAP_PROP_FRAME_HEIGHT | |
| ) | |
| ) | |
| fps = cap.get( | |
| cv2.CAP_PROP_FPS | |
| ) | |
| if fps <= 0: | |
| fps = 30 | |
| process_width = PROCESS_WIDTH | |
| scale_ratio = process_width / width | |
| process_height = int( | |
| height * scale_ratio | |
| ) | |
| scaled_points = [] | |
| for point in brush_points: | |
| scale_x = ( | |
| process_width / | |
| point["canvasWidth"] | |
| ) | |
| scale_y = ( | |
| process_height / | |
| point["canvasHeight"] | |
| ) | |
| scaled_points.append({ | |
| "x": | |
| int( | |
| point["x"] * | |
| scale_x | |
| ), | |
| "y": | |
| int( | |
| point["y"] * | |
| scale_y | |
| ), | |
| "size": | |
| int( | |
| point["size"] * | |
| scale_x | |
| ) | |
| }) | |
| fourcc = cv2.VideoWriter_fourcc( | |
| *"mp4v" | |
| ) | |
| temp_video = os.path.join( | |
| session_dir, | |
| "temp.mp4" | |
| ) | |
| writer = cv2.VideoWriter( | |
| temp_video, | |
| fourcc, | |
| fps, | |
| ( | |
| process_width, | |
| process_height | |
| ) | |
| ) | |
| while True: | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| frame = cv2.resize( | |
| frame, | |
| ( | |
| process_width, | |
| process_height | |
| ) | |
| ) | |
| cleaned = remove_watermark( | |
| frame, | |
| scaled_points | |
| ) | |
| writer.write( | |
| cleaned | |
| ) | |
| cap.release() | |
| writer.release() | |
| subprocess.run([ | |
| "ffmpeg", | |
| "-y", | |
| "-i", | |
| temp_video, | |
| "-i", | |
| input_path, | |
| "-map", | |
| "0:v", | |
| "-map", | |
| "1:a?", | |
| "-c:v", | |
| "copy", | |
| "-c:a", | |
| "aac", | |
| output_path | |
| ]) | |
| mime_type = mimetypes.guess_type( | |
| output_path | |
| )[0] | |
| if mime_type is None: | |
| mime_type = "video/mp4" | |
| return FileResponse( | |
| output_path, | |
| media_type=mime_type, | |
| filename=output_filename, | |
| background=BackgroundTask( | |
| cleanup, | |
| session_dir | |
| ) | |
| ) | |
| # ========================================================= | |
| # MAIN | |
| # ========================================================= | |
| if __name__ == "__main__": | |
| import uvicorn | |
| uvicorn.run( | |
| "app:app", | |
| host="0.0.0.0", | |
| port=7860, | |
| reload=False | |
| ) |