Update app.py

app.py

@@ -9,11 +9,16 @@ 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
 
-#HF_TOKEN = os.getenv("HF_TOKEN")
+# Set your SerpApi key here
+SERPAPI_KEY = "YOUR_SERPAPI_KEY"
 
 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/).
@@ -23,32 +28,29 @@ HEADER = """
 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.
 """
 
-
-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 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
 
 def preprocess(input_image, do_remove_background, foreground_ratio):
     def fill_background(image):
@@ -68,7 +70,6 @@ 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]
@@ -79,23 +80,52 @@ def generate(image):
     mesh.export(mesh_path2.name)
     return mesh_path.name, mesh_path2.name
 
-def run_example(image_pil):
-    preprocessed = preprocess(image_pil, False, 0.9)
-    mesh_name, mesn_name2 = generate(preprocessed)
-    return preprocessed, mesh_name, mesh_name2
+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()
 
 with gr.Blocks() as demo:
     gr.Markdown(HEADER)
     with gr.Row(variant="panel"):
         with gr.Column():
             with gr.Row():
-                input_image = gr.Image(
-                    label="Input Image",
-                    image_mode="RGBA",
-                    sources="upload",
-                    type="pil",
-                    elem_id="content_image",
-                )
+                make_input = gr.Textbox(label="Motorcycle Make", placeholder="Enter motorcycle make")
+                model_input = gr.Textbox(label="Motorcycle Model", placeholder="Enter motorcycle model")
                 processed_image = gr.Image(label="Processed Image", interactive=False)
             with gr.Row():
                 with gr.Group():
@@ -122,27 +152,8 @@ with gr.Blocks() as demo:
                     label="Output Model",
                     interactive=False,
                 )
-    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],
-    )
+    submit.click(fn=run_example, inputs=[make_input, model_input], outputs=[processed_image, output_model, output_model2])
 
 demo.queue(max_size=10)
 demo.launch()
+