# !pip install gradio import gradio as gr import torch from transformers import pipeline from huggingface_hub import InferenceClient from PIL import Image import numpy as np import cv2 # Инициализация моделей # segmentation = pipeline("image-segmentation", model="nvidia/segformer-b5-finetuned-ade-640-640") classification = pipeline("image-classification", model="google/vit-base-patch16-224") upscaling_client = InferenceClient(model="stabilityai/stable-diffusion-x4-upscaler") inpainting_client = InferenceClient(model="stabilityai/stable-diffusion-inpainting") trellis_client = InferenceClient(model="microsoft/TRELLIS") from transformers import OneFormerProcessor, OneFormerForUniversalSegmentation import torch from PIL import Image import numpy as np device = "cuda" if torch.cuda.is_available() else "cpu" processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_ade20k_swin_tiny") model = OneFormerForUniversalSegmentation.from_pretrained("shi-labs/oneformer_ade20k_swin_tiny").to(device) # Функции для обработки изображений from PIL import Image, ImageDraw import numpy as np def segment_image(image): image = Image.fromarray(image) # Изменяем task_input на "panoptic" inputs = processor(image, task_inputs=["panoptic"], return_tensors="pt") with torch.no_grad(): outputs = model(**inputs) # post-process the raw predictions predicted_panoptic_map = processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0] # Extract segment ids and masks segmentation_map = predicted_panoptic_map["segmentation"].cpu().numpy() segments_info = predicted_panoptic_map["segments_info"] # Create cropped masks cropped_masks_with_labels = [] label_counts = {} for segment in segments_info: mask = (segmentation_map == segment["id"]).astype(np.uint8) * 255 cropped_image = cv2.bitwise_and(np.array(image), np.array(image), mask=mask) label = model.config.id2label[segment["label_id"]] # Check if label already exists if label in label_counts: label_counts[label] += 1 else: label_counts[label] = 1 label = f"{label}_{label_counts[label] - 1}" # Append _0, _1, etc. cropped_masks_with_labels.append((cropped_image, label)) return cropped_masks_with_labels def merge_segments_by_labels(gallery_images, labels_input): """ Объединяет сегменты из галереи изображений в одно изображение, основываясь на введенных пользователем метках. Args: gallery_images: Список изображений сегментов (кортежи (изображение, метка)). labels_input: Строка с метками, разделенными точкой с запятой. Returns: Список изображений, где выбранные сегменты объединены в одно. """ # 1. Разделяем входную строку с метками на список labels_to_merge = [label.strip() for label in labels_input.split(";")] # 2. Создаем пустое изображение для объединения merged_image = None # 3. Инициализируем список для хранения индексов объединенных сегментов merged_indices = [] # 4. Проходим по всем изображениям в галерее for i, (image_path, label) in enumerate(gallery_images): # 5. Если метка сегмента в списке меток для объединения if label in labels_to_merge: image = cv2.imread(image_path) # 6. Если это первый сегмент для объединения if merged_image is None: # 7. Создаем копию изображения как основу для объединения merged_image = image.copy() else: # 8. Объединяем текущее изображение с merged_image # Используем cv2.add для наложения изображений, # предполагая, что сегменты не перекрываются merged_image = cv2.add(merged_image, image) # 9. Добавляем индекс объединенного сегмента в список merged_indices.append(i) # 10. Если сегменты были объединены if merged_image is not None: # 11. Создаем новый список изображений, удаляя объединенные сегменты # и добавляя объединенное изображение с новой меткой new_gallery_images = [ item for i, item in enumerate(gallery_images) if i not in merged_indices ] new_name = labels_to_merge[0] new_gallery_images.append((merged_image, new_name)) return new_gallery_images else: # 12. Если не было меток для объединения, возвращаем исходный список return gallery_images def set_client_for_session(request: gr.Request): x_ip_token = request.headers['x-ip-token'] # The "JeffreyXiang/TRELLIS" space is a ZeroGPU space return Client("JeffreyXiang/TRELLIS", headers={"X-IP-Token": x_ip_token}) def generate_3d_model(client, segment_output, segment_name): for i, (image_path, label) in enumerate(segment_output): if label == segment_name: result = client.predict( image=handle_file(image_path), multiimages=[], seed=0, ss_guidance_strength=7.5, ss_sampling_steps=12, slat_guidance_strength=3, slat_sampling_steps=12, multiimage_algo="stochastic", api_name="/image_to_3d" ) break print(result) return result["video"] 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 with gr.Blocks() as demo: client = gr.State() gr.Markdown("# Анализ и редактирование помещений") with gr.Tab("Сканирование"): with gr.Row(): with gr.Column(scale=5): image_input = gr.Image() segment_button = gr.Button("Сегментировать") with gr.Column(scale=5): segment_output = gr.Gallery() 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(): with gr.Column(scale=5): trellis_input = gr.Textbox(label="Имя сегмента для 3D") trellis_button = gr.Button("3D Trellis") with gr.Column(scale=5): trellis_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300) trellis_button.click(generate_3d_model, inputs=[client, 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.load(set_client_for_session, None, client) demo.launch(debug=True, show_error=True)