Update app.py

app.py

@@ -2,14 +2,16 @@ import logging
 import os
 import tempfile
 import time
+
 import gradio as gr
 import numpy as np
 import rembg
 import torch
 from PIL import Image
 from functools import partial
+
 from tsr.system import TSR
-from tsr.utils importremove_background, resize_foreground, to_gradio_3d_orientation
+from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation
 
 # HF_TOKEN = os.getenv("HF_TOKEN")
 
@@ -22,77 +24,86 @@ else:
     device = "cpu"
 
 d = os.environ.get("DEVICE", None)
-if d is not None:
-    device = d
+if d!= None:
+    device = d 
 
 model = TSR.from_pretrained(
     "stabilityai/TripoSR",
-    config_name="config",
+    config_name="config.yaml",
     weight_name="model.ckpt",
 #    token=HF_TOKEN
 )
 model.renderer.set_chunk_size(131072)
-
 model.to(device)
 
-rembg_session = rembg.remb_session()
+rembg_session = rembg.new_session()
+
 
 def check_input_image(input_image):
     if input_image is None:
-        raise Gr.Error("No image uploaded!")
+        raise gr.Error("No image uploaded!")
+
 
 def preprocess(input_image, do_remove_background, foreground_ratio):
     def fill_background(image):
         image = np.array(image).astype(np.float32) / 255.0
-        image = image[:, :, :3] * image[:, :, 3:4] + 0.5 * (1 - image[:, :, 3:4])
-        image = Image.fromarray((image * 255.0).astype(np.uint8)
+        image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
+        image = Image.fromarray((image * 255.0).astype(np.uint8))
+        return image
 
     if do_remove_background:
         image = input_image.convert("RGB")
         image = remove_background(image, rembg_session)
         image = resize_foreground(image, foreground_ratio)
         image = fill_background(image)
-
+    else:
+        image = input_image
+        if image.mode == "RGBA":
+            image = fill_background(image)
     return image
 
+
 def generate(image):
     scene_codes = model(image, device=device)
     mesh = model.extract_mesh(scene_codes)[0]
     mesh = to_gradio_3d_orientation(mesh)
-    mesh_path = tempfile.NamedTemporaryFile(suffix=".obj", delete):
+    mesh_path = tempfile.NamedTemporaryFile(suffix=".obj", delete=False)
     mesh_path2 = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
-    mesh_path2 = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
-
+    mesh.export(mesh_path.name)
+    mesh.export(mesh_path2.name)
     return mesh_path.name, mesh_path2.name
 
-def run_example(image_pил):
-    preprocessed = preprocess(image_pил, False, 0.9)
-    mesh_name, mes_path2.name = generate(image_pил)
 
-    return preprocessed, mesh_name, mesh_name2.name
+def run_example(image_pil):
+    preprocessed = preprocess(image_pil, False, 0.9)
+    mesh_name, mesn_name2 = generate(preprocessed)
+    return preprocessed, mesh_name, mesn_name2
 
-class CustomTheme(Gr.themes.Base):
+
+class CustomTheme(gr.themes.Base):
     def __init__(self):
-        self.primary_h = "#191a1e"
+        super().__init__()
+        self.primary_hue = "#191a1e"
         self.background_fill_primary = "#191a1e"
         self.background_fill_secondary = "#191a1e"
         self.background_fill_tertiary = "#191a1e"
         self.text_color_primary = "#FFFFFF"
         self.text_color_secondary = "#FFFFFF"
-        self.text_color_tertiар = "#FFFFFF"
+        self.text_color_tertiary = "#FFFFFF"
         self.input_background_fill = "#191a1e"
         self.input_text_color = "#FFFFFF"
 
+
 css = """
-/* Скрываем нижление */
+/* Скрываем нижний колонтитул */
 footer {
-    display: none;
+    visibility: hidden;
     height: 0;
     margin: 0;
     padding: 0;
     overflow: hidden;
 }
-@import url('https://fonts.googleapis.com/css2?family=Poppins:w сем/css_code = """
+@import url('https://fonts.googleapis.com/css2?familyPoppins:wght@400;500;700&display=swap');
 /* Применяем шрифты */
 body, input, button, textarea, select,.gr-button {
     font-family: 'Poppins', sans-serif;
@@ -105,13 +116,12 @@ h1, h2, h3, h4, h5, h6 {
     font-weight: 700;
     color: #FFFFFF;
 }
-/* Стиль для текстовых пололей */
+/* Стиль для текстовых полей и кнопок */
 input[type="text"], textarea {
-    background-color: #191a11e!important;
+    background-color: #191a1e!important;
     color: #FFFFFF;
     border: 1px solid #FFFFFF;
 }
-
 /* Цвет кнопки Generate */
 .generate-button {
     background-color: #5271FF!important;
@@ -121,8 +131,7 @@ input[type="text"], textarea {
 }
 .generate-button:hover {
     background-color: #405BBF!important; /* Цвет при наведении */
-}
-
+/* Выравнивание элементов */
 .drop-image-container {
     display: flex;
     flex-direction: column;
@@ -140,75 +149,80 @@ input[type="text"], textarea {
     margin-left: auto;
     margin-right: auto;
 }
+"""
 
-with Gr.ColumnTheme() as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown("**Input Image**", elem_classes="prompt-text")
-
-        with gr.Row():
-            input_image = gr.Image(
-                label="",
-                image_mode="RGBA",
-                sources="upload",
-                type="pil",
-                elem_id="content_image",
-                width=500,  /* Увелинена ширина */
-            processed_image = gr image(label="Processed Image", interactive=False)
-
-        with gr.Row():
-            submitовуй разия:
-                preprocessed = generate(input_image, inputs=[input_image, do_remove_background, foreground_ratio])
-                outputs=[processed_image]
-
-    submit = gr.Button(
-        "Generate",
-        scale=0,
-        variant="primary",
-        elem_classes="generate-button",
-    )
-
-    with gr.Tab("obj"):
-        output_model = gr.Model3D(
-            label="Output Model",
-            interactive=False,
-        )
-
-    with gr.tab("glb"):
-        output_model2 = gr Model3D(
-            label="Output Model",
-            interactive=False,
-        )
-
-with gr.Column(elem_id="col-container"):
-    gr.Markdown("**Input Image**", elem_classes="prompt-text")
-
-    with gr.Row():
-        input_image = gr.Image(
-            label="",
-            image_mode="RGBA",
-            sources="upload",
-            type="pil",
-            elem_id="content_image",
-            width=500,  /* Увелинена ширина */
-            processed_image = gr.Image(label="Processed Image", interactive=False)
-
-    with gr.Row():
-        submit = gr.Button(
-            "Generate",
-            scale=0,
-            variant="primary",
-            elem_classes="generate-button",
-        )
-
-    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],
-    )
+with gr.Blocks(theme=CustomTheme(), css=css) as demo:
+    with gr.Column(elem_id="main-container"):
+        # Заголовок
+        gr.Markdown("**Upload and Process Your Image**", elem_classes="prompt-text")
+
+        # Блок загрузки изображения и настроек
+        with gr.Row(elem_id="upload-and-settings"):
+            # Левый блок с загрузкой изображения
+            with gr.Column(scale=3, elem_id="image-upload-container"):
+                input_image = gr.Image(
+                    label="Upload Image",
+                    image_mode="RGBA",
+                    sources="upload",
+                    type="pil",
+                    elem_id="content_image",
+                    width=500,  # Увеличена ширина входного изображения
+                )
+            # Правый блок с настройками
+            with gr.Column(scale=2, elem_id="settings-container"):
+                do_remove_background = gr.Checkbox(
+                    label="Remove Background",
+                    value=True,
+                )
+                foreground_ratio = gr.Slider(
+                    label="Foreground Ratio",
+                    minimum=0.5,
+                    maximum=1.0,
+                    value=0.85,
+                    step=0.05,
+                )
+
+        # Блок с обработанным изображением
+        with gr.Row(elem_id="processed-image-container"):
+            processed_image = gr.Image(
+                label="Processed Image",
+                interactive=False,
+                elem_id="processed_image",
+            )
+
+        # Кнопка генерации
+        with gr.Row(elem_id="generate-button-container"):
+            submit = gr.Button(
+                "Generate",
+                scale=0,
+                variant="primary",
+                elem_classes="generate-button",
+            )
+
+        # Результаты генерации
+        with gr.Row(elem_id="output-model-container"):
+            with gr.Tab("OBJ Model"):
+                output_model = gr.Model3D(
+                    label="Generated OBJ Model",
+                    interactive=False,
+                )
+            with gr.Tab("GLB Model"):
+                output_model2 = gr.Model3D(
+                    label="Generated GLB Model",
+                    interactive=False,
+                )
+
+# Привязка событий
+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()