Spaces:

mobisoft
/

face-swap-app

Running

App Files Files Community

mobisoft commited on Apr 16

Commit

f8a4ad9

verified ·

1 Parent(s): 459af64

Update app.py

Browse files

Files changed (1) hide show

app.py +120 -83

app.py CHANGED Viewed

@@ -6,6 +6,7 @@ import numpy as np
 import insightface
 import concurrent.futures
 import traceback
 from fastapi import FastAPI, UploadFile, File, HTTPException, Form
 from fastapi.responses import HTMLResponse, StreamingResponse
@@ -31,7 +32,7 @@ TASKS = {}
 executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
 # ============================================================
-# LOAD MODELS
 # ============================================================
 face_app = insightface.app.FaceAnalysis(name="buffalo_l")
 face_app.prepare(ctx_id=-1, det_size=(640, 640))
@@ -48,7 +49,6 @@ def decode_image(file_bytes):
     if img is None and HEIC_SUPPORTED:
         try:
-            from PIL import Image
             import io
             pil_img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
             img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
@@ -56,7 +56,7 @@ def decode_image(file_bytes):
             pass
     if img is None:
-        raise ValueError("Unsupported image format")
     return img
@@ -65,9 +65,8 @@ def compress_and_resize(file_bytes):
     img = decode_image(file_bytes)
     size_mb = len(file_bytes) / (1024 * 1024)
     if size_mb > MAX_FILE_MB:
-        img = cv2.resize(img, None, fx=0.6, fy=0.6, interpolation=cv2.INTER_AREA)
     h, w = img.shape[:2]
     if max(h, w) > MAX_DIM:
@@ -84,19 +83,37 @@ def enhance(img):
 def cleanup():
     now = time.time()
-    remove = []
-    for k, v in TASKS.items():
         if "time" in v and now - v["time"] > CLEANUP_TIME:
             try:
                 if "result" in v:
                     os.remove(v["result"])
             except:
                 pass
-            remove.append(k)
-    for k in remove:
-        TASKS.pop(k, None)
 # ============================================================
 # WORKER
@@ -123,7 +140,7 @@ def run_task(tid, src_bytes, tgt_bytes, filename, face_index):
         name = os.path.splitext(filename)[0]
         out_path = f"/tmp/{name}_{tid}.png"
-        cv2.imwrite(out_path, result, [cv2.IMWRITE_PNG_COMPRESSION, 3])
         TASKS[tid] = {
             "status": "done",
@@ -143,7 +160,7 @@ def run_task(tid, src_bytes, tgt_bytes, filename, face_index):
 app = FastAPI()
 # ============================================================
-# UI PAGE
 # ============================================================
 @app.get("/", response_class=HTMLResponse)
@@ -153,106 +170,112 @@ def home():
 <html>
 <head>
 <meta name="viewport" content="width=device-width, initial-scale=1">
-<title>Face Swap API Test</title>
 <style>
-body{font-family:sans-serif;background:#0f172a;color:white;text-align:center;padding:20px}
 .box{border:2px dashed #444;padding:20px;margin:10px;border-radius:10px;cursor:pointer}
 img{max-width:200px;margin-top:10px;border-radius:8px}
-button{padding:12px 20px;background:#6366f1;color:white;border:none;border-radius:8px}
-.progress{height:10px;background:#333;margin-top:10px;border-radius:10px;overflow:hidden;display:none}
-.bar{height:100%;width:0;background:#22c55e}
 </style>
 </head>
 <body>
-<h2>⚡ Face Swap Test UI</h2>
-<div class="box" onclick="src.click()">Upload Source<input type="file" id="src" hidden></div>
 <img id="p1">
-<div class="box" onclick="tgt.click()">Upload Target<input type="file" id="tgt" hidden></div>
 <img id="p2">
 <br>
-<label>Select Face Index:</label>
-<input type="number" id="faceIndex" value="0" min="0">
-<br><br>
-<button onclick="start()">Start Swap</button>
-<div class="progress" id="progress"><div class="bar" id="bar"></div></div>
-<br>
-<img id="out">
 <br>
-<a id="dl" download="faceswap.png" style="display:none;color:lightgreen">Download PNG</a>
 <script>
-const src=document.getElementById("src");
-const tgt=document.getElementById("tgt");
-src.onchange=()=>p1.src=URL.createObjectURL(src.files[0]);
-tgt.onchange=()=>p2.src=URL.createObjectURL(tgt.files[0]);
-function upload(url,fd){
- return new Promise((res,rej)=>{
-  let xhr=new XMLHttpRequest();
-  xhr.open("POST",url);
-  xhr.upload.onprogress=(e)=>{
-    if(e.lengthComputable){
-      progress.style.display="block";
-      bar.style.width=(e.loaded/e.total*100)+"%";
-    }
-  };
-  xhr.onload=()=>res(JSON.parse(xhr.responseText));
-  xhr.onerror=rej;
-  xhr.send(fd);
- });
-}
 async function start(){
- if(!src.files[0]||!tgt.files[0]) return alert("Upload both");
- let fd=new FormData();
- fd.append("source",src.files[0]);
- fd.append("target",tgt.files[0]);
- fd.append("face_index",document.getElementById("faceIndex").value);
- let data=await upload("/swap",fd);
- poll(data.task_id);
 }
 async function poll(id){
- let r=await fetch("/status/"+id);
- let j=await r.json();
- if(j.status==="done"){
-   let img=await fetch("/result/"+id);
-   let blob=await img.blob();
-   let url=URL.createObjectURL(blob);
-   out.src=url;
-   dl.href=url;
-   dl.style.display="block";
-   bar.style.width="100%";
- }
- else if(j.status==="failed"){
-   alert(j.error);
- }
- else{
-   setTimeout(()=>poll(id),800);
- }
 }
 </script>
@@ -261,7 +284,17 @@ async function poll(id){
 """
 # ============================================================
-# API
 # ============================================================
 @app.post("/swap")
@@ -285,6 +318,9 @@ async def swap(
     return {"task_id": tid}
 @app.get("/status/{tid}")
 def status(tid: str):
@@ -293,6 +329,9 @@ def status(tid: str):
         raise HTTPException(404)
     return TASKS[tid]
 @app.get("/result/{tid}")
 def result(tid: str):
@@ -304,7 +343,5 @@ def result(tid: str):
     return StreamingResponse(
         open(task["result"], "rb"),
         media_type="image/png",
-        headers={
-            "Content-Disposition": f'attachment; filename="{task["filename"]}"'
-        }
     )

 import insightface
 import concurrent.futures
 import traceback
+import base64
 from fastapi import FastAPI, UploadFile, File, HTTPException, Form
 from fastapi.responses import HTMLResponse, StreamingResponse
 executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
 # ============================================================
+# MODELS
 # ============================================================
 face_app = insightface.app.FaceAnalysis(name="buffalo_l")
 face_app.prepare(ctx_id=-1, det_size=(640, 640))
     if img is None and HEIC_SUPPORTED:
         try:
             import io
             pil_img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
             img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
             pass
     if img is None:
+        raise ValueError("Unsupported image")
     return img
     img = decode_image(file_bytes)
     size_mb = len(file_bytes) / (1024 * 1024)
     if size_mb > MAX_FILE_MB:
+        img = cv2.resize(img, None, fx=0.6, fy=0.6)
     h, w = img.shape[:2]
     if max(h, w) > MAX_DIM:
 def cleanup():
     now = time.time()
+    for k in list(TASKS.keys()):
+        v = TASKS[k]
         if "time" in v and now - v["time"] > CLEANUP_TIME:
             try:
                 if "result" in v:
                     os.remove(v["result"])
             except:
                 pass
+            TASKS.pop(k, None)
+# ============================================================
+# FACE DETECTION API
+# ============================================================
+def extract_faces(img):
+    faces = face_app.get(img)
+    results = []
+    for i, f in enumerate(faces):
+        x1, y1, x2, y2 = map(int, f.bbox)
+        crop = img[y1:y2, x1:x2]
+        _, buf = cv2.imencode(".jpg", crop)
+        b64 = base64.b64encode(buf).decode()
+        results.append({
+            "index": i,
+            "image": f"data:image/jpeg;base64,{b64}"
+        })
+    return results
 # ============================================================
 # WORKER
         name = os.path.splitext(filename)[0]
         out_path = f"/tmp/{name}_{tid}.png"
+        cv2.imwrite(out_path, result)
         TASKS[tid] = {
             "status": "done",
 app = FastAPI()
 # ============================================================
+# UI
 # ============================================================
 @app.get("/", response_class=HTMLResponse)
 <html>
 <head>
 <meta name="viewport" content="width=device-width, initial-scale=1">
+<title>AI Face Swap</title>
 <style>
+body{font-family:sans-serif;background:#0f172a;color:white;text-align:center}
 .box{border:2px dashed #444;padding:20px;margin:10px;border-radius:10px;cursor:pointer}
 img{max-width:200px;margin-top:10px;border-radius:8px}
+.faces{display:flex;gap:10px;justify-content:center;margin-top:10px}
+.face{border:2px solid transparent;cursor:pointer}
+.face img{width:60px;height:60px}
+.face.active{border-color:#22c55e}
+.loader{display:none;margin-top:10px}
+.spin{border:4px solid #333;border-top:4px solid #22c55e;border-radius:50%;width:30px;height:30px;animation:spin 1s linear infinite}
+@keyframes spin{100%{transform:rotate(360deg)}}
 </style>
 </head>
 <body>
+<h2>⚡ AI Face Swap</h2>
+<input type="file" id="src"><br>
 <img id="p1">
+<div id="faces" class="faces"></div>
+<input type="file" id="tgt"><br>
 <img id="p2">
 <br>
+<button onclick="start()">Swap</button>
+<div class="loader" id="loader">
+<div class="spin"></div>
+</div>
 <br>
+<img id="out"><br>
+<a id="dl" download="faceswap.png" style="display:none;color:lightgreen">Download</a>
 <script>
+let selectedFace=0;
+src.onchange=async()=>{
+  p1.src=URL.createObjectURL(src.files[0]);
+  let fd=new FormData();
+  fd.append("file",src.files[0]);
+  let res=await fetch("/detect-faces",{method:"POST",body:fd});
+  let data=await res.json();
+  faces.innerHTML="";
+  data.faces.forEach(f=>{
+    let d=document.createElement("div");
+    d.className="face";
+    d.innerHTML="<img src='"+f.image+"'>";
+    d.onclick=()=>{
+      document.querySelectorAll(".face").forEach(x=>x.classList.remove("active"));
+      d.classList.add("active");
+      selectedFace=f.index;
+    };
+    faces.appendChild(d);
+  });
+  if(faces.firstChild) faces.firstChild.classList.add("active");
+};
+tgt.onchange=()=>p2.src=URL.createObjectURL(tgt.files[0]);
 async function start(){
+  loader.style.display="block";
+  let fd=new FormData();
+  fd.append("source",src.files[0]);
+  fd.append("target",tgt.files[0]);
+  fd.append("face_index",selectedFace);
+  let r=await fetch("/swap",{method:"POST",body:fd});
+  let j=await r.json();
+  poll(j.task_id);
 }
 async function poll(id){
+  let r=await fetch("/status/"+id);
+  let j=await r.json();
+  if(j.status==="done"){
+    let img=await fetch("/result/"+id);
+    let blob=await img.blob();
+    let url=URL.createObjectURL(blob);
+    out.src=url;
+    dl.href=url;
+    dl.style.display="block";
+    loader.style.display="none";
+  }
+  else if(j.status==="failed"){
+    alert(j.error);
+    loader.style.display="none";
+  }
+  else{
+    setTimeout(()=>poll(id),800);
+  }
 }
 </script>
 """
 # ============================================================
+# DETECT FACES
+# ============================================================
+@app.post("/detect-faces")
+async def detect_faces(file: UploadFile = File(...)):
+    img = compress_and_resize(await file.read())
+    faces = extract_faces(img)
+    return {"faces": faces}
+# ============================================================
+# SWAP
 # ============================================================
 @app.post("/swap")
     return {"task_id": tid}
+# ============================================================
+# STATUS
+# ============================================================
 @app.get("/status/{tid}")
 def status(tid: str):
         raise HTTPException(404)
     return TASKS[tid]
+# ============================================================
+# RESULT
+# ============================================================
 @app.get("/result/{tid}")
 def result(tid: str):
     return StreamingResponse(
         open(task["result"], "rb"),
         media_type="image/png",
+        headers={"Content-Disposition": f'attachment; filename="{task["filename"]}"'}
     )