semi-final version

app.py

@@ -7,6 +7,10 @@ from gradio_client import Client, handle_file
 import numpy as np
 import cv2
 import os
+import tempfile
+import io
+import base64
+import requests
 
 # Инициализация моделей
 from transformers import OneFormerProcessor, OneFormerForUniversalSegmentation
@@ -22,7 +26,6 @@ oneFormer_model = OneFormerForUniversalSegmentation.from_pretrained("shi-labs/on
 # inpainting_client = InferenceClient(model="stabilityai/stable-diffusion-inpainting")
 # Функции для обработки изображений
 def segment_image(image):
-    image = Image.fromarray(image)
     inputs = oneFormer_processor(image, task_inputs=["panoptic"], return_tensors="pt")
 
     with torch.no_grad():
@@ -91,7 +94,94 @@ def merge_segments_by_labels(gallery_images, labels_input):
     else:
         return gallery_images
 
+def select_segment(segment_output, segment_name):
+    for i, (image_path, label) in enumerate(segment_output):
+        if label == segment_name:
+            return image_path
+
+#Image edit
 
+def return_image(imageEditor):
+    return imageEditor['composite']
+
+def rembg_client(request: gr.Request):
+    try:
+        client = Client("KenjieDec/RemBG", headers={"X-IP-Token": request.headers['x-ip-token']})
+        print("KenjieDec/RemBG Ip token")
+        return client
+    except:
+        print("KenjieDec/RemBG no token")
+        return Client("KenjieDec/RemBG")
+
+def autocrop_image(imageEditor, border = 0):
+    image = imageEditor['composite']
+    bbox = image.getbbox()
+    image = image.crop(bbox)
+    (width, height) = image.size
+    width += border * 2
+    height += border * 2
+    cropped_image = Image.new("RGBA", (width, height), (0,0,0,0))
+    cropped_image.paste(image, (border, border))
+    return cropped_image
+
+def remove_black_make_transparent(imageEditor):
+    image_pil = imageEditor['composite']
+    if image_pil.mode != "RGBA":
+        image_pil = image_pil.convert("RGBA")
+    image_np = np.array(image_pil)
+    black_pixels_mask = np.all(image_np[:, :, :3] == [0, 0, 0], axis=-1)
+    image_np[black_pixels_mask, 3] = 0
+    transparent_image = Image.fromarray(image_np)
+    return transparent_image
+
+def rembg(imageEditor, request: gr.Request):
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+        imageEditor['composite'].save(temp_file.name)
+        temp_file_path = temp_file.name
+        client = rembg_client(request)
+        result = client.predict(
+            file=handle_file(temp_file_path),
+            mask="Default",
+            model="birefnet-general-lite",
+            x=0,
+            y=0,
+            api_name="/inference"
+        )
+        print(result)
+        return result
+
+def add_transparent_border(imageEditor, border_size=200):
+    image = imageEditor['composite']
+    width, height = image.size
+    new_width = width + 2 * border_size
+    new_height = height + 2 * border_size
+    new_image = Image.new("RGBA", (new_width, new_height), (0, 0, 0, 0))
+    new_image.paste(image, (border_size, border_size))
+    return new_image
+
+def upscale(imageEditor, scale, request: gr.Request): 
+    return upscale_image(imageEditor['composite'], version="v1.4", rescaling_factor=scale)
+
+def upscale_image(image_pil, version="v1.4", rescaling_factor=None):
+    buffered = io.BytesIO()
+    image_pil.save(buffered, format="PNG")  # Save as PNG
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+
+    # Update the data format for PNG
+    data = {"data": [f"data:image/png;base64,{img_str}", version, rescaling_factor]}
+    # Send request to the API
+    response = requests.post("https://nightfury-image-face-upscale-restoration-gfpgan.hf.space/api/predict", json=data)
+    response.raise_for_status()
+    # Get the base64 data from the response
+    base64_data = response.json()["data"][0]
+    base64_data = base64_data.split(",")[1] # remove data:image/png;base64,
+    # Convert base64 back to PIL Image
+    image_bytes = base64.b64decode(base64_data)
+    upscaled_image = Image.open(io.BytesIO(image_bytes))
+    return upscaled_image
+
+
+#3d models
 def hunyuan_client(request: gr.Request):
     try:
         client = Client("tencent/Hunyuan3D-2", headers={"X-IP-Token": request.headers['x-ip-token']})
@@ -110,117 +200,115 @@ def vFusion_client(request: gr.Request):
         print("facebook/VFusion3D no token")
         return Client("facebook/VFusion3D")
 
-def generate_3d_model(segment_output, request: gr.Request):
-    for i, (image_path, label) in enumerate(segment_output):
-        if label == segment_name:
-            client = hunyuan_client(request)
-            result = client.predict(
-                caption="",
-                image=handle_file(image_path),
-                steps=50,
-                guidance_scale=5.5,
-                seed=1234,
-                octree_resolution="256",
-                check_box_rembg=True,
-                api_name="/shape_generation"
-            )
-            print(result)
-            return result[0]
-
-def generate_3d_model_texture(segment_output, segment_name, request: gr.Request):
-    for i, (image_path, label) in enumerate(segment_output):
-        if label == segment_name:
-            client = hunyuan_client(request)
-            result = client.predict(
-                caption="",
-                image=handle_file(image_path),
-                steps=50,
-                guidance_scale=5.5,
-                seed=1234,
-                octree_resolution="256",
-                check_box_rembg=True,
-                api_name="/generation_all"
-            )
-            print(result)
-            return result[1]
-
-def generate_3d_model2(segment_output, segment_name, request: gr.Request):
-    for i, (image_path, label) in enumerate(segment_output):
-        if label == segment_name:
-            client = vFusion_client(request)
-            result = client.predict(
-                image=handle_file(image_path),
-                api_name="/step_1_generate_obj"
-            )
-            print(result)
-            return result[0]
-
-
-# def classify_segments(segments):
-#     # Предполагается, что segments - список изображений сегментов
-#     results = []
-#     for segment in segments:
-#         results.append(classification(segment))
-#     return results  # Вернем список классификаций
-
-# def upscale_segment(segment):
-#     upscaled = upscaling_client.image_to_image(segment)
-#     return upscaled
-
-# def inpaint_image(image, mask, prompt):
-#     inpainted = inpainting_client.text_to_image(prompt, image=image, mask=mask)
-#     return inpainted
-
-
+def generate_3d_model(image_pil, rembg_Hunyuan, request: gr.Request):
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+        image_pil.save(temp_file.name)
+        temp_file_path = temp_file.name
+        client = hunyuan_client(request)
+        result = client.predict(
+            caption="",
+            image=handle_file(temp_file_path),
+            steps=50,
+            guidance_scale=5.5,
+            seed=1234,
+            octree_resolution="256",
+            check_box_rembg=rembg_Hunyuan,
+            api_name="/shape_generation"
+        )
+        print(result)
+        return result[0]
+
+def generate_3d_model_texture(image_pil, rembg_Hunyuan, request: gr.Request):
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+        image_pil.save(temp_file.name)
+        temp_file_path = temp_file.name
+        client = hunyuan_client(request)
+        result = client.predict(
+            caption="",
+            image=handle_file(temp_file_path),
+            steps=50,
+            guidance_scale=5.5,
+            seed=1234,
+            octree_resolution="256",
+            check_box_rembg=rembg_Hunyuan,
+            api_name="/generation_all"
+        )
+        print(result)
+        return result[1]
+
+def generate_3d_model2(image_pil, request: gr.Request):
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+        image_pil.save(temp_file.name)
+        temp_file_path = temp_file.name
+        client = vFusion_client(request)
+        result = client.predict(
+            image=handle_file(temp_file_path),
+            api_name="/step_1_generate_obj"
+        )
+        print(result)
+        return result[0]
 
 
 ########## GRADIO ##########
 
+
 with gr.Blocks() as demo:
     gr.Markdown("# Анализ и редактирование помещений")
-
     with gr.Tab("Сканирование"):
-        with gr.Row():
+        with gr.Row(equal_height=True):
             with gr.Column(scale=5):
-                image_input = gr.Image()
+                image_input = gr.Image(type="pil", label="Исходное изображение", height = 400)
                 segment_button = gr.Button("Сегментировать")
             with gr.Column(scale=5):
-                segment_output = gr.Gallery()
+                segments_output = gr.Gallery(label="Сегменты изображения")
                 merge_segments_input = gr.Textbox(label="Сегменты для объединения (через точку с запятой, например: \"wall_0; tv_0\")")
                 merge_segments_button = gr.Button("Соединить сегменты")
-                merge_segments_button.click(merge_segments_by_labels, inputs=[segment_output, merge_segments_input], outputs=segment_output)
-        with gr.Row():
+                merge_segments_button.click(merge_segments_by_labels, inputs=[segments_output, merge_segments_input], outputs=segments_output)
+                with gr.Row(equal_height=True):
+                    segment_text_input = gr.Textbox(label="Имя сегмента для дальнейшего редактирования")
+                    select_segment_button = gr.Button("Использовать сегмент")
+    with gr.Tab("Редактирование"):
+        with gr.Row(equal_height=True):
             with gr.Column(scale=5):
-                trellis_input = gr.Textbox(label="Имя сегмента для 3D")
-                hunyuan_button = gr.Button("Hunyuan3D-2")
-                hunyuan_button_texture = gr.Button("Hunyuan3D-2 (with texture)")
-                vFusion_button = gr.Button("VFusion3D")
+                segment_input = gr.ImageEditor(type="pil", label="Сегмент для редактирования")
+            with gr.Column(scale=5):
+                crop_button = gr.Button("Обрезать сегмент")
+                with gr.Row(equal_height=True):
+                    upscale_slider = gr.Slider(minimum=1, maximum=5, value=2, step=0.1, label="во сколько раз")
+                    upscale_button = gr.Button("Upscale")
+                rembg_button = gr.Button("Rembg")
+                remove_background_button = gr.Button("Убрать черный задний фон")
+                with gr.Row(equal_height=True):
+                    add_transparent_border_slider = gr.Slider(minimum=10, maximum=500, value=200, step=10, label="в пикселях")
+                    add_transparent_border_button = gr.Button("Добавить прозрачные края")
+                use_button = gr.Button("Использовать сегмент для 3D")
+                    
+    with gr.Tab("Создание 3D"):
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=5):
+                segment_3d_input = gr.Image(type="pil", image_mode="RGBA", label="Сегмент для 3D", height = 600)
+                rembg_Hunyuan = gr.Checkbox(label="Hunyuan3D-2 rembg Enabled", info="Включить rembg для Hunyuan3D-2?")
+                hunyuan_button = gr.Button("Hunyuan3D-2 (no texture) [ZeroGPU = 100s]")
+                hunyuan_button_texture = gr.Button("Hunyuan3D-2 (with texture) [ZeroGPU = 150s]")
+                vFusion_button = gr.Button("VFusion3D [если у вас совсем все грустно по ZeroGPU]")
             with gr.Column(scale=5):
                 trellis_output = gr.Model3D(label="3D Model")
-                
-                hunyuan_button.click(generate_3d_model, inputs=[segment_output, trellis_input], outputs=trellis_output)
-                hunyuan_button_texture.click(generate_3d_model_texture, inputs=[segment_output, trellis_input], outputs=trellis_output)
-                vFusion_button.click(generate_3d_model2, inputs=[segment_output, trellis_input], outputs=trellis_output)
-
-        segment_button.click(segment_image, inputs=image_input, outputs=segment_output)
-        # segment_button.click(segment_full_image, inputs=image_input, outputs=segment_output)
-
-    # with gr.Tab("Редактирование"):
-    #     segment_input = gr.Image()
-    #     upscale_output = gr.Image()
-    #     upscale_button = gr.Button("Upscale")
-    #     upscale_button.click(upscale_segment, inputs=segment_input, outputs=upscale_output)
-
-    #     mask_input = gr.Image()
-    #     prompt_input = gr.Textbox()
-    #     inpaint_output = gr.Image()
-    #     inpaint_button = gr.Button("Inpaint")
-    #     inpaint_button.click(inpaint_image, inputs=[segment_input, mask_input, prompt_input], outputs=inpaint_output)
-
-    # with gr.Tab("Создание 3D моделей"):
-    #     segment_input_3d = gr.Image()
-    #     model_output = gr.File()
-    #     model_button = gr.Button("Создать 3D модель")
-    #     model_button.click(generate_3d_model, inputs=segment_input_3d, outputs=model_output)
-
-demo.launch(debug=True, show_error=True)
+
+    #tab1
+    segment_button.click(segment_image, inputs=image_input, outputs=segments_output)
+    select_segment_button.click(select_segment, inputs=[segments_output, segment_text_input], outputs=segment_input)
+    
+    #tab2
+    crop_button.click(autocrop_image, inputs=segment_input, outputs=segment_input)
+    upscale_button.click(upscale, inputs=[segment_input, upscale_slider], outputs=segment_input)
+    rembg_button.click(rembg, inputs=segment_input, outputs=segment_input)
+    remove_background_button.click(remove_black_make_transparent, inputs=segment_input, outputs=segment_input)
+    add_transparent_border_button.click(add_transparent_border, inputs=[segment_input, add_transparent_border_slider], outputs=segment_input)
+    use_button.click(return_image, inputs=segment_input, outputs=segment_3d_input)
+
+    #3d buttons
+    hunyuan_button.click(generate_3d_model, inputs=[segment_3d_input, rembg_Hunyuan], outputs=trellis_output)
+    hunyuan_button_texture.click(generate_3d_model_texture, inputs=[segment_3d_input, rembg_Hunyuan], outputs=trellis_output)
+    vFusion_button.click(generate_3d_model2, inputs=segment_3d_input, outputs=trellis_output)
+
+demo.launch(debug=True, show_error=True)