Update app.py

app.py

@@ -6,116 +6,102 @@ import numpy as np
 import insightface
 import concurrent.futures
 import traceback
-import requests
 
 from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.responses import HTMLResponse, StreamingResponse
 
+# HEIC SUPPORT
+from PIL import Image
+import pillow_heif
+pillow_heif.register_heif_opener()
+
 # ============================================================
 # CONFIG
 # ============================================================
 MAX_FILE_MB = 10
 MAX_DIM = 640
-UPSCALE_FACTOR = 2
+UPSCALE_SIZE = 1024   # HD output
 MAX_WORKERS = 3
 CLEANUP_TIME = 300
 
 TASKS = {}
 executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
 
-BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-
 # ============================================================
 # LOAD MODELS
 # ============================================================
 face_app = insightface.app.FaceAnalysis(name="buffalo_l")
 face_app.prepare(ctx_id=-1, det_size=(640, 640))
 
-swapper = insightface.model_zoo.get_model("inswapper_128.onnx", root=BASE_DIR)
+swapper = insightface.model_zoo.get_model("inswapper_128.onnx", root=".")
 
 # ============================================================
-# GFPGAN LITE (FACE ENHANCER)
+# IMAGE DECODING (ALL FORMATS SUPPORT)
 # ============================================================
-try:
-    import onnxruntime as ort
-    ONNX_AVAILABLE = True
-except:
-    ONNX_AVAILABLE = False
 
-gfpgan_session = None
-GFPGAN_URL = "https://huggingface.co/ai-forever/GFPGAN-Lite/resolve/main/gfpgan_lite.onnx"
-
-def load_gfpgan():
-    global gfpgan_session
-    if not ONNX_AVAILABLE:
-        return
+def read_image(file_bytes):
+    try:
+        # try OpenCV first
+        img = cv2.imdecode(np.frombuffer(file_bytes, np.uint8), cv2.IMREAD_COLOR)
+        if img is not None:
+            return img
 
-    path = os.path.join(BASE_DIR, "gfpgan_lite.onnx")
+        # fallback for HEIC / unsupported
+        pil_img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
+        return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
 
-    if not os.path.exists(path):
-        print("Downloading GFPGAN...")
-        r = requests.get(GFPGAN_URL)
-        open(path, "wb").write(r.content)
+    except Exception:
+        raise ValueError("Unsupported image format")
 
-    gfpgan_session = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
-    print("GFPGAN Loaded")
 
 # ============================================================
-# IMAGE HELPERS
+# OPTIMIZATION
 # ============================================================
 
-def decode_image(b):
-    img = cv2.imdecode(np.frombuffer(b, np.uint8), cv2.IMREAD_COLOR)
-    if img is None:
-        raise ValueError("Invalid image")
-    return img
-
-
-def compress_resize(b):
-    img = decode_image(b)
+def compress_and_resize(file_bytes):
+    img = read_image(file_bytes)
 
-    size_mb = len(b) / (1024 * 1024)
+    size_mb = len(file_bytes) / (1024 * 1024)
 
+    # compress large images
     if size_mb > MAX_FILE_MB:
         img = cv2.resize(img, None, fx=0.6, fy=0.6, interpolation=cv2.INTER_AREA)
 
+    # resize for faster inference
     h, w = img.shape[:2]
     if max(h, w) > MAX_DIM:
         scale = MAX_DIM / max(h, w)
-        img = cv2.resize(img, (int(w * scale), int(h * scale)))
+        img = cv2.resize(img, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_LINEAR)
 
     return img
 
 
+# ============================================================
+# HD UPSCALE (FAST CPU FRIENDLY)
+# ============================================================
+
 def upscale_hd(img):
     h, w = img.shape[:2]
-    img = cv2.resize(img, (w * UPSCALE_FACTOR, h * UPSCALE_FACTOR), interpolation=cv2.INTER_CUBIC)
 
-    # sharpening
-    kernel = np.array([[0,-1,0],[-1,5,-1],[0,-1,0]])
-    img = cv2.filter2D(img, -1, kernel)
+    # upscale to HD target
+    scale = UPSCALE_SIZE / max(h, w)
 
-    return img
-
-
-def enhance_face(img):
-    if gfpgan_session is None:
-        return img
+    img = cv2.resize(
+        img,
+        (int(w * scale), int(h * scale)),
+        interpolation=cv2.INTER_CUBIC
+    )
 
-    try:
-        inp = cv2.resize(img, (512, 512))
-        inp = inp.astype(np.float32) / 255.0
-        inp = np.transpose(inp, (2, 0, 1))[None]
+    # sharpen for better detail
+    blur = cv2.GaussianBlur(img, (0, 0), 1.0)
+    img = cv2.addWeighted(img, 1.3, blur, -0.3, 0)
 
-        out = gfpgan_session.run(None, {"input": inp})[0]
-        out = np.transpose(out[0], (1, 2, 0))
-        out = (out * 255).clip(0,255).astype(np.uint8)
+    return img
 
-        out = cv2.resize(out, (img.shape[1], img.shape[0]))
-        return out
-    except:
-        return img
 
+# ============================================================
+# CLEANUP
+# ============================================================
 
 def cleanup():
     now = time.time()
@@ -133,16 +119,17 @@ def cleanup():
     for k in remove:
         TASKS.pop(k, None)
 
+
 # ============================================================
 # WORKER
 # ============================================================
 
-def run_task(tid, src_b, tgt_b):
+def run_task(tid, src_bytes, tgt_bytes):
     TASKS[tid]["status"] = "processing"
 
     try:
-        src = compress_resize(src_b)
-        tgt = compress_resize(tgt_b)
+        src = compress_and_resize(src_bytes)
+        tgt = compress_and_resize(tgt_bytes)
 
         s_faces = face_app.get(src)
         t_faces = face_app.get(tgt)
@@ -152,12 +139,11 @@ def run_task(tid, src_b, tgt_b):
 
         result = swapper.get(tgt, t_faces[0], s_faces[0], paste_back=True)
 
-        # ===== HD UPSCALE =====
-        result = enhance_face(result)
+        # HD UPSCALE
         result = upscale_hd(result)
 
-        out_path = f"/tmp/{tid}.webp"
-        cv2.imwrite(out_path, result, [cv2.IMWRITE_WEBP_QUALITY, 90])
+        out_path = f"/tmp/{tid}.jpg"
+        cv2.imwrite(out_path, result, [cv2.IMWRITE_JPEG_QUALITY, 95])
 
         TASKS[tid] = {
             "status": "done",
@@ -169,15 +155,114 @@ def run_task(tid, src_b, tgt_b):
         TASKS[tid] = {"status": "failed", "error": str(e)}
         print(traceback.format_exc())
 
+
 # ============================================================
 # FASTAPI
 # ============================================================
 
 app = FastAPI()
 
+# ============================================================
+# UI
+# ============================================================
+
 @app.get("/", response_class=HTMLResponse)
 def home():
-    return "<h2>FaceSwap Running (HD Enabled)</h2>"
+    return """
+<!DOCTYPE html>
+<html>
+<head>
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>HD Face Swap</title>
+
+<style>
+body{background:#0f172a;color:white;text-align:center;font-family:sans-serif}
+.container{max-width:900px;margin:auto;padding:20px}
+img{width:100%;max-height:260px;object-fit:contain;border-radius:10px}
+button{padding:12px 18px;margin:10px;background:#6366f1;color:white;border:none;border-radius:8px}
+.download{display:none;background:#10b981}
+</style>
+</head>
+
+<body>
+
+<div class="container">
+
+<h2>🔥 HD Face Swap (iOS Ready)</h2>
+
+<input type="file" id="src"><br><br>
+<input type="file" id="tgt"><br><br>
+
+<img id="p1"><br>
+<img id="p2"><br>
+
+<button onclick="start()">Upload & Swap</button>
+
+<p id="status"></p>
+
+<img id="out"><br>
+<a id="dl" class="download" download="faceswap_hd.jpg">Download HD</a>
+
+</div>
+
+<script>
+const src = document.getElementById("src");
+const tgt = document.getElementById("tgt");
+const dl = document.getElementById("dl");
+
+src.onchange = ()=> p1.src = URL.createObjectURL(src.files[0]);
+tgt.onchange = ()=> p2.src = URL.createObjectURL(tgt.files[0]);
+
+async function start(){
+  if(!src.files[0] || !tgt.files[0]){
+    alert("Upload both images");
+    return;
+  }
+
+  let fd = new FormData();
+  fd.append("source", src.files[0]);
+  fd.append("target", tgt.files[0]);
+
+  status.innerText = "Processing...";
+
+  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();
+
+  status.innerText = j.status;
+
+  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 = "inline-block";
+
+    status.innerText = "Done ✅";
+  } else if(j.status==="failed"){
+    status.innerText = j.error;
+  } else {
+    setTimeout(()=>poll(id), 1000);
+  }
+}
+</script>
+
+</body>
+</html>
+"""
+
+# ============================================================
+# API (iOS FRIENDLY)
+# ============================================================
 
 @app.post("/swap")
 async def swap(source: UploadFile = File(...), target: UploadFile = File(...)):
@@ -198,8 +283,10 @@ async def swap(source: UploadFile = File(...), target: UploadFile = File(...)):
 @app.get("/status/{tid}")
 def status(tid: str):
     cleanup()
+
     if tid not in TASKS:
         raise HTTPException(404)
+
     return TASKS[tid]
 
 
@@ -210,12 +297,4 @@ def result(tid: str):
     if not task or task["status"] != "done":
         raise HTTPException(404)
 
-    return StreamingResponse(open(task["result"], "rb"), media_type="image/webp")
-
-# ============================================================
-# INIT
-# ============================================================
-
-print("Loading models...")
-load_gfpgan()
-print("Ready 🚀")
+    return StreamingResponse(open(task["result"], "rb"), media_type="image/jpeg")