Update app.py

app.py

@@ -14,6 +14,13 @@ from PIL import Image
 import sf3d.utils as sf3d_utils
 from sf3d.system import SF3D
 
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import FileResponse
+
+import datetime
+
+app = FastAPI()
+
 rembg_session = rembg.new_session()
 
 COND_WIDTH = 512
@@ -167,189 +174,27 @@ def show_mask_img(input_image: Image) -> Image:
     return Image.fromarray(new_img.astype(np.uint8), mode="RGB")
 
 
-def run_button(run_btn, input_image, background_state, foreground_ratio):
-    if run_btn == "Run":
-        glb_file: str = run_model(background_state)
-
-        return (
-            gr.update(),
-            gr.update(),
-            gr.update(),
-            gr.update(),
-            gr.update(value=glb_file, visible=True),
-            gr.update(visible=True),
-        )
-    elif run_btn == "Remove Background":
-        rem_removed = remove_background(input_image)
-
-        sqr_crop = square_crop(rem_removed)
-        fr_res = resize_foreground(sqr_crop, foreground_ratio)
-
-        return (
-            gr.update(value="Run", visible=True),
-            sqr_crop,
-            fr_res,
-            gr.update(value=show_mask_img(fr_res), visible=True),
-            gr.update(value=None, visible=False),
-            gr.update(visible=False),
-        )
-
-
-def requires_bg_remove(image, fr):
-    if image is None:
-        return (
-            gr.update(visible=False, value="Run"),
-            None,
-            None,
-            gr.update(value=None, visible=False),
-            gr.update(visible=False),
-            gr.update(visible=False),
-        )
-    alpha_channel = np.array(image.getchannel("A"))
-    min_alpha = alpha_channel.min()
-
-    if min_alpha == 0:
-        print("Already has alpha")
-        sqr_crop = square_crop(image)
-        fr_res = resize_foreground(sqr_crop, fr)
-        return (
-            gr.update(value="Run", visible=True),
-            sqr_crop,
-            fr_res,
-            gr.update(value=show_mask_img(fr_res), visible=True),
-            gr.update(visible=False),
-            gr.update(visible=False),
-        )
-    return (
-        gr.update(value="Remove Background", visible=True),
-        None,
-        None,
-        gr.update(value=None, visible=False),
-        gr.update(visible=False),
-        gr.update(visible=False),
-    )
-
 
-def update_foreground_ratio(img_proc, fr):
-    foreground_res = resize_foreground(img_proc, fr)
-    return (
-        foreground_res,
-        gr.update(value=show_mask_img(foreground_res)),
-    )
+def upload_file_to_s3(file_path, bucket_name, object_name=None):
+    
+    s3_client.upload_file(file_path, bucket_name, object_name)
 
+    return True
 
-with gr.Blocks() as demo:
-    img_proc_state = gr.State()
-    background_remove_state = gr.State()
-    gr.Markdown("""
-    # SF3D: Stable Fast 3D Mesh Reconstruction with UV-unwrapping and Illumination Disentanglement
-
-    **SF3D** is a state-of-the-art method for 3D mesh reconstruction from a single image.
-    This demo allows you to upload an image and generate a 3D mesh model from it.
-
-    **Tips**
-    1. If the image already has an alpha channel, you can skip the background removal step.
-    2. You can adjust the foreground ratio to control the size of the foreground object. This can influence the shape
-    3. You can upload your own HDR environment map to light the 3D model.
-    """)
-    with gr.Row(variant="panel"):
-        with gr.Column():
-            with gr.Row():
-                input_img = gr.Image(
-                    type="pil", label="Input Image", sources="upload", image_mode="RGBA"
-                )
-                preview_removal = gr.Image(
-                    label="Preview Background Removal",
-                    type="pil",
-                    image_mode="RGB",
-                    interactive=False,
-                    visible=False,
-                )
-
-            foreground_ratio = gr.Slider(
-                label="Foreground Ratio",
-                minimum=0.5,
-                maximum=1.0,
-                value=0.85,
-                step=0.05,
-            )
-
-            foreground_ratio.change(
-                update_foreground_ratio,
-                inputs=[img_proc_state, foreground_ratio],
-                outputs=[background_remove_state, preview_removal],
-            )
-
-            run_btn = gr.Button("Run", variant="primary", visible=False)
-
-        with gr.Column():
-            output_3d = LitModel3D(
-                label="3D Model",
-                visible=False,
-                clear_color=[0.0, 0.0, 0.0, 0.0],
-                tonemapping="aces",
-                contrast=1.0,
-                scale=1.0,
-            )
-            with gr.Column(visible=False, scale=1.0) as hdr_row:
-                gr.Markdown("""## HDR Environment Map
-
-                Select an HDR environment map to light the 3D model. You can also upload your own HDR environment maps.
-                """)
-
-                with gr.Row():
-                    hdr_illumination_file = gr.File(
-                        label="HDR Env Map", file_types=[".hdr"], file_count="single"
-                    )
-                    example_hdris = [
-                        os.path.join("demo_files/hdri", f)
-                        for f in os.listdir("demo_files/hdri")
-                    ]
-                    hdr_illumination_example = gr.Examples(
-                        examples=example_hdris,
-                        inputs=hdr_illumination_file,
-                    )
-
-                    hdr_illumination_file.change(
-                        lambda x: gr.update(env_map=x.name if x is not None else None),
-                        inputs=hdr_illumination_file,
-                        outputs=[output_3d],
-                    )
-
-    examples = gr.Examples(
-        examples=example_files,
-        inputs=input_img,
-    )
+    
+@app.post("/process-image/")
+async def process_image(file: UploadFile = File(...), foreground_ratio: float = 0.85):
+    input_image = Image.open(file.file).convert("RGBA")
+    rem_removed = remove_background(input_image)
+    sqr_crop = square_crop(rem_removed)
+    fr_res = resize_foreground(sqr_crop, foreground_ratio)
+    glb_file = run_model(fr_res)
 
-    input_img.change(
-        requires_bg_remove,
-        inputs=[input_img, foreground_ratio],
-        outputs=[
-            run_btn,
-            img_proc_state,
-            background_remove_state,
-            preview_removal,
-            output_3d,
-            hdr_row,
-        ],
-    )
 
-    run_btn.click(
-        run_button,
-        inputs=[
-            run_btn,
-            input_img,
-            background_remove_state,
-            foreground_ratio,
-        ],
-        outputs=[
-            run_btn,
-            img_proc_state,
-            background_remove_state,
-            preview_removal,
-            output_3d,
-            hdr_row,
-        ],
-    )
+    timestamp = datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')
+    object_name = f'object_{timestamp}.glb'
 
-demo.launch()
+    if upload_file_to_s3(glb_file, 'framebucket3d',object_name):
+        return {
+            "glb_path": f"https://framebucket3d.s3.amazonaws.com/{object_name_2}"
+        }