Update app.py

app.py

@@ -9,48 +9,45 @@ import rembg
 import torch
 from PIL import Image
 from functools import partial
-from serpapi import GoogleSearch
-import requests
-from io import BytesIO
-import matplotlib.pyplot as plt
 
 from tsr.system import TSR
 from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation
 
-# Set your SerpApi key here
-SERPAPI_KEY = "3a786d94adb1d9739bb3a877b05dae35d231917f02ad89f74adfa014b567af3f"
+#HF_TOKEN = os.getenv("HF_TOKEN")
 
 HEADER = """
 **TripoSR** is a state-of-the-art open-source model for **fast** feedforward 3D reconstruction from a single image, developed in collaboration between [Tripo AI](https://www.tripo3d.ai/) and [Stability AI](https://stability.ai/).
-
 **Tips:**
 1. If you find the result is unsatisfied, please try to change the foreground ratio. It might improve the results.
 2. Please disable "Remove Background" option only if your input image is RGBA with transparent background, image contents are centered and occupy more than 70% of image width or height.
 """
 
-def get_motorcycle_image(make, model):
-    params = {
-        "api_key": SERPAPI_KEY,
-        "engine": "google",
-        "q": f"{make} {model} motorcycle product photo",
-        "tbm": "isch"
-    }
-
-    search = GoogleSearch(params)
-    results = search.get_dict()
-    if "images_results" in results:
-        first_image = results["images_results"][0]
-        image_url = first_image.get("original")
-        if image_url:
-            image_response = requests.get(image_url)
-            image = Image.open(BytesIO(image_response.content))
-            return image
-        else:
-            print("Image URL not found in results.")
-            return None
-    else:
-        print("No image results found.")
-        return None
+
+if torch.cuda.is_available():
+    device = "cuda:0"
+else:
+    device = "cpu"
+
+d = os.environ.get("DEVICE", None)
+if d != None:
+    device = d 
+
+model = TSR.from_pretrained(
+    "stabilityai/TripoSR",
+    config_name="config.yaml",
+    weight_name="model.ckpt",
+#    token=HF_TOKEN
+)
+model.renderer.set_chunk_size(131072)
+model.to(device)
+
+rembg_session = rembg.new_session()
+
+
+def check_input_image(input_image):
+    if input_image is None:
+        raise gr.Error("No image uploaded!")
+
 
 def preprocess(input_image, do_remove_background, foreground_ratio):
     def fill_background(image):
@@ -70,6 +67,7 @@ def preprocess(input_image, do_remove_background, foreground_ratio):
             image = fill_background(image)
     return image
 
+
 def generate(image):
     scene_codes = model(image, device=device)
     mesh = model.extract_mesh(scene_codes)[0]
@@ -80,52 +78,23 @@ def generate(image):
     mesh.export(mesh_path2.name)
     return mesh_path.name, mesh_path2.name
 
-def run_example(make, model):
-    image = get_motorcycle_image(make, model)
-    if image:
-        # Save the image
-        input_image_path = '/content/motorcycle.jpg'
-        image.save(input_image_path)
-
-        # Load the image
-        img = Image.open(input_image_path)
-        output_image_path = '/content/motorcyclebg.png'
-        img_no_bg = rembg_remove(img)
-        img_no_bg.save(output_image_path)
-
-        # Preprocess and generate 3D model
-        preprocessed = preprocess(img_no_bg, False, 0.9)
-        mesh_name, mesh_name2 = generate(preprocessed)
-        return preprocessed, mesh_name, mesh_name2
-    else:
-        raise gr.Error("Image could not be fetched.")
-
-if torch.cuda.is_available():
-    device = "cuda:0"
-else:
-    device = "cpu"
-
-d = os.environ.get("DEVICE", None)
-if d != None:
-    device = d 
-
-model = TSR.from_pretrained(
-    "stabilityai/TripoSR",
-    config_name="config.yaml",
-    weight_name="model.ckpt",
-)
-model.renderer.set_chunk_size(131072)
-model.to(device)
-
-rembg_session = rembg.new_session()
+def run_example(image_pil):
+    preprocessed = preprocess(image_pil, False, 0.9)
+    mesh_name, mesn_name2 = generate(preprocessed)
+    return preprocessed, mesh_name, mesh_name2
 
 with gr.Blocks() as demo:
     gr.Markdown(HEADER)
     with gr.Row(variant="panel"):
         with gr.Column():
             with gr.Row():
-                make_input = gr.Textbox(label="Motorcycle Make", placeholder="Enter motorcycle make")
-                model_input = gr.Textbox(label="Motorcycle Model", placeholder="Enter motorcycle model")
+                input_image = gr.Image(
+                    label="Input Image",
+                    image_mode="RGBA",
+                    sources="upload",
+                    type="pil",
+                    elem_id="content_image",
+                )
                 processed_image = gr.Image(label="Processed Image", interactive=False)
             with gr.Row():
                 with gr.Group():
@@ -152,8 +121,27 @@ with gr.Blocks() as demo:
                     label="Output Model",
                     interactive=False,
                 )
-    submit.click(fn=run_example, inputs=[make_input, model_input], outputs=[processed_image, output_model, output_model2])
+    with gr.Row(variant="panel"):
+        gr.Examples(
+            examples=[
+                os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
+            ],
+            inputs=[input_image],
+            outputs=[processed_image, output_model, output_model2],
+            #cache_examples=True,
+            fn=partial(run_example),
+            label="Examples",
+            examples_per_page=20
+        )
+    submit.click(fn=check_input_image, inputs=[input_image]).success(
+        fn=preprocess,
+        inputs=[input_image, do_remove_background, foreground_ratio],
+        outputs=[processed_image],
+    ).success(
+        fn=generate,
+        inputs=[processed_image],
+        outputs=[output_model, output_model2],
+    )
 
 demo.queue(max_size=10)
-demo.launch()
-
+demo.launch()