import os os.system("pip uninstall -y huggingface-hub") os.system("pip install huggingface-hub==0.25.2") import gradio as gr from PIL import Image import numpy as np import cv2 from zipfile import ZipFile # Функция обработки изображения import gradio as gr from PIL import Image import numpy as np import os import shutil import yaml from pathlib import Path from fabric_diffusion import FabricDiffusionPipeline from delete_bg import remove_background import capture from pytorch_lightning import Trainer import random import torch ZIP_FOLDER = "./ZIPS" os.makedirs(ZIP_FOLDER, exist_ok=True) def zip_folder(folder_path, output_zip): with ZipFile(output_zip, 'w') as zipf: for root, dirs, files in os.walk(folder_path): for file in files: file_path = os.path.join(root, file) zipf.write(file_path, os.path.relpath(file_path, folder_path)) def set_deterministic(seed=42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def load_config(config_path): with open(config_path, 'r') as file: return yaml.safe_load(file) def run_flatten_texture(pipeline, input_image_path, output_path, n_samples=3): os.makedirs(output_path, exist_ok=True) texture_name = os.path.splitext(os.path.basename(input_image_path))[0] texture_patch = pipeline.load_patch_data(input_image_path) gen_imgs = pipeline.flatten_texture(texture_patch, n_samples=n_samples) for i, gen_img in enumerate(gen_imgs): gen_img.save(os.path.join(output_path, f'{texture_name}_gen_{i}.png')) def organize_images_into_structure(source_folder, new_folder): os.makedirs(new_folder, exist_ok=True) for file_name in os.listdir(source_folder): source_file = os.path.join(source_folder, file_name) if os.path.isfile(source_file) and file_name.lower().endswith(('.png', '.jpg', '.jpeg')): folder_name = os.path.splitext(file_name)[0] subfolder_path = os.path.join(new_folder, folder_name, "outputs") os.makedirs(subfolder_path, exist_ok=True) destination_file = os.path.join(subfolder_path, file_name) shutil.copy(source_file, destination_file) # Create a directory for saving if it doesn't exist if not os.path.exists("saved_images"): os.makedirs("saved_images") def decode_rgba(image: Image.Image) -> Image.Image: image_array = np.array(image) alpha_channel = image_array[:, :, 3] coords = np.argwhere(alpha_channel == 255) y_min, x_min = coords.min(axis=0) y_max, x_max = coords.max(axis=0) + 1 # Include the max boundary cropped_image_array = image_array[y_min:y_max, x_min:x_max] cropped_rgb_array = cropped_image_array[:, :, :3] cropped_rgb_image = Image.fromarray(cropped_rgb_array, "RGB") return cropped_rgb_image def process_image(orig_image, image_data): """Processes and saves the original and cropped images.""" # Get the original image. Convert to PIL Image if needed # Get the cropped image (composite). Convert to PIL Image if needed cropped_image = image_data['composite'] # print(type(cropped_image), cropped_image.size, np.array(crop_image).shape) # Generate unique filenames using timestamps to avoid overwriting import datetime timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S") original_filename = os.path.join("saved_images", f"original_{timestamp}.png") cropped_filename = os.path.join("saved_images", f"cropped_{timestamp}.png") orig_image.save(original_filename) decoded_cropped_image = decode_rgba(cropped_image) decoded_cropped_image.save(cropped_filename) return orig_image, original_filename, decoded_cropped_image, cropped_filename, timestamp def web_main(orig_image, image_data): orig_image, original_filename, cropped_image, cropped_filename, timestamp = process_image(orig_image, image_data) set_deterministic(seed=42) config = load_config("config_demo.yaml") device = config["hyperparameters"]["device"] texture_checkpoint = config["hyperparameters"]["fb_checkpoint"] print_checkpoint = config["hyperparameters"].get("print_checkpoint", None) input_image = config["hyperparameters"]["input_image"] save_fd_dir = config["hyperparameters"]["save_fd_dir"] save_mp_dir = config["hyperparameters"]["save_mp_dir"] # save_mp_dir = timestamp n_samples = config["hyperparameters"]["n_samples"] pipeline = FabricDiffusionPipeline(device, texture_checkpoint, print_checkpoint=print_checkpoint) os.makedirs(save_fd_dir, exist_ok=True) run_flatten_texture(pipeline, cropped_filename, output_path=save_fd_dir, n_samples=n_samples) organize_images_into_structure(save_fd_dir, save_mp_dir) data = capture.get_data(predict_dir=Path(save_mp_dir), predict_ds='sd') module = capture.get_inference_module(pt=config["hyperparameters"]["checkpoint_name"]) decomp = Trainer(default_root_dir=Path(save_mp_dir), accelerator='gpu', devices=1, precision=16) decomp.predict(module, data) folder = f"cropped_{timestamp}_gen_0" # for folder in os.listdir(save_mp_dir): folder_path = os.path.join(save_mp_dir, folder) if os.path.isdir(folder_path): target_path = os.path.join(folder_path, "weights", "mask", "an_object_with_azertyuiop_texture", "checkpoint-800", "outputs") if os.path.exists(target_path): for file_name in os.listdir(target_path): file_path = os.path.join(target_path, file_name) if os.path.isfile(file_path): shutil.move(file_path, folder_path) print(f"FOLDER: {folder_path}") print(os.path.exists(original_filename)) shutil.copyfile( original_filename, os.path.join(folder_path, "outputs", original_filename.split("/")[-1]) ) shutil.copyfile( cropped_filename, os.path.join(folder_path, "outputs", cropped_filename.split("/")[-1]) ) zip_folder( folder_path, os.path.join(ZIP_FOLDER, f"{folder_path.split('/')[-1]}.zip") ) return orig_image, cropped_image, os.path.join(ZIP_FOLDER, f"{folder_path.split('/')[-1]}.zip") def remove_background_main(image): import datetime timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S") if not os.path.exists("saved_images"): os.makedirs("saved_images") original_filename = os.path.join("saved_images", f"original_{timestamp}.png") image.save(original_filename) result_img = remove_background(original_filename) output_path = os.path.join("saved_images", f"nobg_{timestamp}.png") cv2.imwrite(output_path, cv2.cvtColor(result_img, cv2.COLOR_RGBA2BGRA)) folder_name = f"nobg_{timestamp}" folder_path = os.path.join("output_images", folder_name) os.makedirs(os.path.join(folder_path, "outputs"), exist_ok=True) shutil.copyfile( original_filename, os.path.join(folder_path, "outputs", os.path.basename(original_filename)) ) shutil.copyfile( output_path, os.path.join(folder_path, "outputs", os.path.basename(output_path)) ) result_pil = Image.fromarray(cv2.cvtColor(result_img, cv2.COLOR_RGBA2BGRA)) return result_pil with gr.Blocks() as demo: with gr.Tabs() as tabs: # PBR generator # with gr.TabItem("PBR generator") as pbr_tab: # with gr.Row(): # orig_image = gr.Image(label="Orig image", type="pil") # image_editor = gr.ImageEditor(type="pil", crop_size="1:1", label="Edit Image: Crop the desired element") # with gr.Row(): # output_image = gr.Image(label="Result") # crop_image = gr.Image(label="Crop_image") # zip_file = gr.File(label="Download Zip File") # process_button = gr.Button("Crop Element") # Print extraction with gr.TabItem("Print extraction") as print_tab: with gr.Row(): orig_image_print = gr.Image(label="Orig image (cropped print)", type="pil") with gr.Row(): output_image_print = gr.Image(label="Result") process_print_button = gr.Button("Extract print") # process_button.click( # web_main, # inputs=[orig_image, image_editor], # outputs=[output_image, crop_image, zip_file] # ) process_print_button.click( remove_background_main, inputs=[orig_image_print], outputs=[output_image_print] ) if __name__ == "__main__": demo.launch()