Update app.py

app.py

@@ -11,6 +11,9 @@ import subprocess
 from glob import glob
 import requests
 from huggingface_hub import snapshot_download
+import io
+from fastapi import UploadFile, File, HTTPException
+from fastapi.responses import FileResponse
 
 # Download models
 os.makedirs("ckpts", exist_ok=True)
@@ -37,49 +40,35 @@ images_examples = [
 ]
 
 def remove_background(input_pil, remove_bg):
-    
-    # Create a temporary folder for downloaded and processed images
     temp_dir = tempfile.mkdtemp()
     unique_id = str(uuid.uuid4())
     image_path = os.path.join(temp_dir, f'input_image_{unique_id}.png')  
-       
     try:
-        # Check if input_url is already a PIL Image
         if isinstance(input_pil, Image.Image):
             image = input_pil
         else:
-            # Otherwise, assume it's a file path and open it
             image = Image.open(input_pil)
-        
-        # Flip the image horizontally
         image = image.transpose(Image.FLIP_LEFT_RIGHT)
-        
-        # Save the resized image
         image.save(image_path)
     except Exception as e:
         shutil.rmtree(temp_dir)
         raise gr.Error(f"Error downloading or saving the image: {str(e)}")
 
     if remove_bg is True:
-        # Run background removal
         removed_bg_path = os.path.join(temp_dir, f'output_image_rmbg_{unique_id}.png')
         try:            
             img = Image.open(image_path)    
             result = remove(img)
             result.save(removed_bg_path)
-
-            # Remove the input image to keep the temp directory clean
             os.remove(image_path)
         except Exception as e:
             shutil.rmtree(temp_dir)
             raise gr.Error(f"Error removing background: {str(e)}")
-
         return removed_bg_path, temp_dir
     else: 
         return image_path, temp_dir
     
 def run_inference(temp_dir, removed_bg_path):
-    # Define the inference configuration
     inference_config = "configs/inference-768-6view.yaml"
     pretrained_model = "./ckpts"
     crop_size = 740
@@ -88,7 +77,6 @@ def run_inference(temp_dir, removed_bg_path):
     save_mode = "rgb"
 
     try:
-        # Run the inference command
         subprocess.run(
             [
                 "python", "inference.py",
@@ -103,19 +91,12 @@ def run_inference(temp_dir, removed_bg_path):
             ],
             check=True
         )
-
-        
-        # Retrieve the file name without the extension
         removed_bg_file_name = os.path.splitext(os.path.basename(removed_bg_path))[0]
-
-        # List objects in the "out" folder
         out_folder_path = "out"
         out_folder_objects = os.listdir(out_folder_path)
         print(f"Objects in '{out_folder_path}':")
         for obj in out_folder_objects:
             print(f" - {obj}")
-        
-        # List objects in the "out/{removed_bg_file_name}" folder
         specific_out_folder_path = os.path.join(out_folder_path, removed_bg_file_name)
         if os.path.exists(specific_out_folder_path) and os.path.isdir(specific_out_folder_path):
             specific_out_folder_objects = os.listdir(specific_out_folder_path)
@@ -124,36 +105,23 @@ def run_inference(temp_dir, removed_bg_path):
                 print(f" - {obj}")
         else:
             print(f"\nThe folder '{specific_out_folder_path}' does not exist.")
-        
         output_video = glob(os.path.join(f"out/{removed_bg_file_name}", "*.mp4"))
         output_objects = glob(os.path.join(f"out/{removed_bg_file_name}", "*.obj"))
         return output_video, output_objects
-    
     except subprocess.CalledProcessError as e:
         return f"Error during inference: {str(e)}"
 
 def process_image(input_pil, remove_bg, progress=gr.Progress(track_tqdm=True)):
-
     torch.cuda.empty_cache()
-    
-    # Remove background
     result = remove_background(input_pil, remove_bg)
-    
     if isinstance(result, str) and result.startswith("Error"):
-        raise gr.Error(f"{result}")  # Return the error message if something went wrong
-
-    removed_bg_path, temp_dir = result  # Unpack only if successful
-
-    # Run inference
+        raise gr.Error(f"{result}")
+    removed_bg_path, temp_dir = result
     output_video, output_objects = run_inference(temp_dir, removed_bg_path)
-
     if isinstance(output_video, str) and output_video.startswith("Error"):
         shutil.rmtree(temp_dir)
-        raise gr.Error(f"{output_video}")   # Return the error message if inference failed
-
-    
-    shutil.rmtree(temp_dir)  # Cleanup temporary folder
-    print(output_video)
+        raise gr.Error(f"{output_video}")
+    shutil.rmtree(temp_dir)
     torch.cuda.empty_cache()
     return output_video[0], output_objects[0], output_objects[1]
 
@@ -166,6 +134,7 @@ div#video-out-elm{
     height: 323px;
 }
 """
+
 def gradio_interface():
     with gr.Blocks(css=css) as app:
         with gr.Column(elem_id="col-container"):
@@ -178,48 +147,55 @@ def gradio_interface():
                 <a href="https://penghtyx.github.io/PSHuman/">
                     <img src='https://img.shields.io/badge/Project-Page-green'>
                 </a>
-    			<a href="https://arxiv.org/pdf/2409.10141">
+                <a href="https://arxiv.org/pdf/2409.10141">
                     <img src='https://img.shields.io/badge/ArXiv-Paper-red'>
                 </a>
                 <a href="https://huggingface.co/spaces/fffiloni/PSHuman?duplicate=true">
-    				<img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space">
-    			</a>
-    			<a href="https://huggingface.co/fffiloni">
-    				<img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-me-on-HF-sm-dark.svg" alt="Follow me on HF">
-    			</a>
+                    <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space">
+                </a>
+                <a href="https://huggingface.co/fffiloni">
+                    <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-me-on-HF-sm-dark.svg" alt="Follow me on HF">
+                </a>
             </div>
             """)
             with gr.Group():
                 with gr.Row():  
                     with gr.Column(scale=2):
-                        
-                        input_image = gr.Image(
-                            label="Image input", 
-                            type="pil",
-                            image_mode="RGBA",
-                            height=480
-                        )
-        
+                        input_image = gr.Image(label="Image input", type="pil", image_mode="RGBA", height=480)
                         remove_bg = gr.Checkbox(label="Need to remove BG ?", value=False)
-                    
                         submit_button = gr.Button("Process")
-                    
                     with gr.Column(scale=4):
                         output_video= gr.Video(label="Output Video", elem_id="video-out-elm")
                         with gr.Row():
                             output_object_mesh = gr.Model3D(label=".OBJ Mesh", height=240)
                             output_object_color = gr.Model3D(label=".OBJ colored", height=240)
-    
             gr.Examples(
                 examples = examples_folder,
                 inputs = [input_image],
                 examples_per_page = 11
             )
-
         submit_button.click(process_image, inputs=[input_image, remove_bg], outputs=[output_video, output_object_mesh, output_object_color])
-
     return app
 
-# Launch the Gradio app
-app = gradio_interface()
-app.launch(show_api=False, show_error=True, ssr_mode=False)
+if __name__ == "__main__":
+    gradio_app = gradio_interface()
+    fastapi_app = gradio_app.app
+
+    @fastapi_app.post("/api/3d-reconstruct")
+    async def reconstruct(
+        image_file: UploadFile = File(...),
+        remove_bg: bool = False
+    ):
+        try:
+            contents = await image_file.read()
+            pil_image = Image.open(io.BytesIO(contents)).convert("RGBA")
+            video_path, mesh_path, colored_path = process_image(pil_image, remove_bg)
+            return FileResponse(
+                colored_path,
+                media_type="application/octet-stream",
+                filename=os.path.basename(colored_path)
+            )
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
+
+    gradio_app.launch(show_api=False, show_error=True, ssr_mode=False)