Update app.py

app.py

@@ -1,5 +1,7 @@
 import os
 from io import BytesIO
+import uuid
+import threading
 
 import cv2
 import numpy as np
@@ -21,6 +23,19 @@ app = FastAPI(
     version="1.0.0"
 )
 
+# ====================================================
+# Job Storage
+# ====================================================
+
+jobs = {}
+results = {}
+
+def update_progress(job_id, value, message):
+    jobs[job_id] = {
+        "progress": value,
+        "message": message
+    }
+
 # ====================================================
 # Load Model Once
 # ====================================================
@@ -66,10 +81,13 @@ face_helper = FaceRestoreHelper(
 # Face Enhancement Function
 # ====================================================
 
-def enhance_face(image_rgb):
-
+def enhance_face(image_rgb, job_id=None):
+    
     face_helper.clean_all()
-
+    
+    if job_id:
+        update_progress(job_id, 5, "Preparing image")
+    
     image_bgr = cv2.cvtColor(
         image_rgb,
         cv2.COLOR_RGB2BGR
@@ -77,15 +95,36 @@ def enhance_face(image_rgb):
 
     face_helper.read_image(image_bgr)
 
+    if job_id:
+        update_progress(job_id, 20, "Detecting faces")
+    
     face_helper.get_face_landmarks_5(
         only_center_face=False,
         resize=640,
         eye_dist_threshold=5
     )
 
+    if job_id:
+        update_progress(job_id, 35, "Aligning faces")
+    
     face_helper.align_warp_face()
-
-    for cropped_face in face_helper.cropped_faces:
+    
+    total_faces = len(face_helper.cropped_faces)
+    
+    if total_faces == 0:
+        if job_id:
+            update_progress(job_id, 100, "No faces found")
+        return image_rgb
+
+    for idx, cropped_face in enumerate(face_helper.cropped_faces):
+        
+        if job_id:
+            progress = 40 + int((idx / total_faces) * 50)
+            update_progress(
+                job_id,
+                progress,
+                f"Enhancing face {idx+1}/{total_faces}"
+            )
 
         cropped_face_t = img2tensor(
             cropped_face / 255.0,
@@ -135,6 +174,9 @@ def enhance_face(image_rgb):
             cropped_face
         )
 
+    if job_id:
+        update_progress(job_id, 95, "Finalizing image")
+    
     face_helper.get_inverse_affine(None)
 
     restored_img = face_helper.paste_faces_to_input_image()
@@ -143,10 +185,46 @@ def enhance_face(image_rgb):
         restored_img,
         cv2.COLOR_BGR2RGB
     )
+    
+    if job_id:
+        update_progress(job_id, 100, "Completed")
 
     return restored_img
 
 
+# ====================================================
+# Background Worker
+# ====================================================
+
+def process_job(job_id, image):
+    try:
+        result = enhance_face(image, job_id)
+        
+        result = cv2.cvtColor(
+            result,
+            cv2.COLOR_RGB2BGR
+        )
+        
+        _, buffer = cv2.imencode(
+            ".png",
+            result
+        )
+        
+        results[job_id] = buffer.tobytes()
+        
+        update_progress(
+            job_id,
+            100,
+            "Completed"
+        )
+        
+    except Exception as e:
+        update_progress(
+            job_id,
+            -1,
+            str(e)
+        )
+
 # ====================================================
 # UI
 # ====================================================
@@ -189,6 +267,22 @@ HTML_PAGE = """
         h1{
             text-align:center;
         }
+        
+        .progress-container {
+            width:100%;
+            border:1px solid #ccc;
+            height:30px;
+            margin:10px 0;
+            border-radius:4px;
+            overflow:hidden;
+        }
+        
+        .progress-bar {
+            height:100%;
+            width:0%;
+            background:#4CAF50;
+            transition: width 0.3s ease;
+        }
     </style>
 
 </head>
@@ -206,6 +300,10 @@ Enhance Face
 
 <p id="status"></p>
 
+<div class="progress-container">
+    <div id="bar" class="progress-bar"></div>
+</div>
+
 <div class="container">
 
     <div class="card">
@@ -247,7 +345,10 @@ async function enhance(){
     }
 
     document.getElementById("status")
-    .innerHTML = "Processing...";
+    .innerHTML = "Starting...";
+    
+    document.getElementById("bar")
+    .style.width = "0%";
 
     const formData = new FormData();
 
@@ -256,22 +357,61 @@ async function enhance(){
         file
     );
 
-    const response = await fetch(
-        "/convert",
-        {
-            method:"POST",
-            body:formData
-        }
+    const start =
+        await fetch("/convert", {
+            method: "POST",
+            body: formData
+        });
+
+    const data =
+        await start.json();
+
+    const jobId =
+        data.job_id;
+
+    const timer = setInterval(
+        async () => {
+
+            const res =
+                await fetch(
+                    "/progress/" + jobId
+                );
+
+            const progress =
+                await res.json();
+
+            document.getElementById(
+                "status"
+            ).innerHTML =
+                progress.message;
+
+            document.getElementById(
+                "bar"
+            ).style.width =
+                progress.progress + "%";
+
+            if(progress.progress >= 100){
+
+                clearInterval(timer);
+
+                document
+                    .getElementById("outputPreview")
+                    .src =
+                    "/result/" +
+                    jobId +
+                    "?t=" +
+                    Date.now();
+            }
+            
+            if(progress.progress < 0){
+                clearInterval(timer);
+                document.getElementById("status")
+                .innerHTML = "Error: " + progress.message;
+            }
+
+        },
+        500
     );
-
-    const blob =
-    await response.blob();
-
-    document.getElementById("outputPreview")
-    .src = URL.createObjectURL(blob);
-
-    document.getElementById("status")
-    .innerHTML = "Done";
 }
 
 </script>
@@ -286,49 +426,67 @@ async def home():
 
 
 # ====================================================
-# API
+# API Endpoints
 # ====================================================
 
 @app.post("/convert")
-async def convert(
-    file: UploadFile = File(...)
-):
-
+async def convert(file: UploadFile = File(...)):
     contents = await file.read()
-
+    
     np_img = np.frombuffer(
         contents,
         np.uint8
     )
-
+    
     image = cv2.imdecode(
         np_img,
         cv2.IMREAD_COLOR
     )
-
+    
     image = cv2.cvtColor(
         image,
         cv2.COLOR_BGR2RGB
     )
+    
+    job_id = str(uuid.uuid4())
+    
+    jobs[job_id] = {
+        "progress": 0,
+        "message": "Queued"
+    }
+    
+    threading.Thread(
+        target=process_job,
+        args=(job_id, image),
+        daemon=True
+    ).start()
+    
+    return {
+        "job_id": job_id
+    }
 
-    result = enhance_face(image)
-
-    result = cv2.cvtColor(
-        result,
-        cv2.COLOR_RGB2BGR
-    )
-
-    _, buffer = cv2.imencode(
-        ".png",
-        result
+@app.get("/progress/{job_id}")
+async def progress(job_id: str):
+    return jobs.get(
+        job_id,
+        {
+            "progress": 0,
+            "message": "Unknown job"
+        }
     )
 
+@app.get("/result/{job_id}")
+async def result(job_id: str):
+    if job_id not in results:
+        return {
+            "status": "processing"
+        }
+    
     return StreamingResponse(
-        BytesIO(buffer.tobytes()),
+        BytesIO(results[job_id]),
         media_type="image/png"
     )
 
-
 # ====================================================
 # Run
 # ====================================================