Spaces:
Sleeping
Sleeping
| # import gradio as gr | |
| # import torch | |
| # import os | |
| # import tempfile | |
| # import shutil | |
| # from PIL import Image | |
| # import numpy as np | |
| # from pathlib import Path | |
| # import sys | |
| # import copy | |
| # # --- Import logic from your project --- | |
| # from options.test_options import TestOptions | |
| # from data import create_dataset | |
| # from models import create_model | |
| # try: | |
| # from best_ldr import compute_metrics_for_images, score_records | |
| # except ImportError: | |
| # raise ImportError("Could not import from best_ldr.py. Make sure the file is in the same directory as app.py.") | |
| # print("--- Initializing LDR-to-HDR Model (this may take a moment) ---") | |
| # # --- Global Setup: Load the CycleGAN model once when the app starts --- | |
| # # We need to satisfy the parser's requirement for a dataroot at startup | |
| # if '--dataroot' not in sys.argv: | |
| # sys.argv.extend(['--dataroot', './dummy_dataroot_for_init']) | |
| # # Load the base options | |
| # opt = TestOptions().parse() | |
| # # Manually override settings for our model | |
| # opt.name = 'ldr2hdr_cyclegan_728' | |
| # opt.model = 'test' | |
| # opt.netG = 'resnet_9blocks' | |
| # opt.norm = 'instance' | |
| # opt.no_dropout = True | |
| # opt.checkpoints_dir = './checkpoints' | |
| # opt.gpu_ids = [0] if torch.cuda.is_available() else [] | |
| # opt.device = torch.device('cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids else torch.device('cpu') | |
| # # Create the model using these options | |
| # model = create_model(opt) | |
| # model.setup(opt) | |
| # model.eval() | |
| # print("--- Model Loaded Successfully ---") | |
| # # --- Helper Function for Inference --- | |
| # def run_inference(model, image_path, process_options): | |
| # """ | |
| # A reusable function to run the model with specific preprocessing options. | |
| # """ | |
| # # Deep copy the base options to avoid modifying the global state | |
| # local_opt = copy.deepcopy(opt) | |
| # # Apply the specific settings for this run | |
| # for key, value in process_options.items(): | |
| # setattr(local_opt, key, value) | |
| # with tempfile.TemporaryDirectory() as temp_dir: | |
| # shutil.copy(image_path, temp_dir) | |
| # local_opt.dataroot = temp_dir | |
| # local_opt.num_test = 1 | |
| # dataset = create_dataset(local_opt) | |
| # for i, data in enumerate(dataset): | |
| # model.set_input(data) | |
| # model.test() | |
| # visuals = model.get_current_visuals() | |
| # for label, image_tensor in visuals.items(): | |
| # if label == 'fake': | |
| # image_numpy = (np.transpose(image_tensor.cpu().float().numpy()[0], (1, 2, 0)) + 1) / 2.0 * 255.0 | |
| # return Image.fromarray(image_numpy.astype(np.uint8)) | |
| # # --- The Main Gradio Processing Function --- | |
| # def process_images_and_display(list_of_temp_files): | |
| # """ | |
| # The main workflow: select best LDR, then run two inference modes. | |
| # """ | |
| # if not list_of_temp_files: | |
| # raise gr.Error("Please upload your bracketed LDR images.") | |
| # if len(list_of_temp_files) < 2: | |
| # gr.Warning("For best results, upload at least 2 bracketed LDR images.") | |
| # uploaded_filepaths = [Path(f.name) for f in list_of_temp_files] | |
| # try: | |
| # # --- Step 1: Select the Best LDR --- | |
| # print(f"Analyzing {len(uploaded_filepaths)} uploaded images...") | |
| # weights = {"clipped": 0.35, "coverage": 0.25, "exposure": 0.15, "sharpness": 0.15, "noise": 0.10} | |
| # records = compute_metrics_for_images(uploaded_filepaths, resize_max=1024) | |
| # scored_records = score_records(records, weights) | |
| # if not scored_records: | |
| # raise gr.Error("Could not read or score any of the uploaded images.") | |
| # best_ldr_record = scored_records[0] | |
| # best_ldr_path = best_ldr_record['path'] | |
| # print(f"Best LDR selected: {os.path.basename(best_ldr_path)} (Score: {best_ldr_record['score']:.4f})") | |
| # chosen_ldr_image = Image.open(best_ldr_path).convert("RGB") | |
| # # --- Step 2: Run Inference in Both Modes --- | |
| # # Mode A: High-Quality Crop (at model's native resolution) | |
| # print("Running Mode A: High-Quality Crop...") | |
| # crop_options = { | |
| # 'preprocess': 'resize_and_crop', | |
| # 'load_size': 728, | |
| # 'crop_size': 728 | |
| # } | |
| # hdr_cropped = run_inference(model, best_ldr_path, crop_options) | |
| # print("Mode A successful.") | |
| # # Mode B: Full Image (at a higher resolution) | |
| # print("Running Mode B: Full Image (High-Res Scaled)...") | |
| # scale_options = { | |
| # 'preprocess': 'scale_width', | |
| # 'load_size': 1024, # <-- THIS IS THE CHANGE FOR HIGHER RESOLUTION | |
| # 'crop_size': 728 # This value is ignored by scale_width but needs to be present | |
| # } | |
| # hdr_scaled = run_inference(model, best_ldr_path, scale_options) | |
| # print("Mode B successful.") | |
| # # Return all the images to update the UI | |
| # return uploaded_filepaths, chosen_ldr_image, hdr_cropped, hdr_scaled | |
| # except Exception as e: | |
| # print(f"An error occurred: {e}") | |
| # raise gr.Error(f"An error occurred during processing: {e}") | |
| # # --- Create and Launch the Gradio Interface --- | |
| # with gr.Blocks(theme=gr.themes.Monochrome(), css="footer {display: none !important}") as demo: | |
| # gr.Markdown("# LDR Bracketing to HDR Converter") | |
| # gr.Markdown("Upload a set of bracketed LDR images. The app will automatically select the best one and convert it to HDR using two different methods for comparison.") | |
| # with gr.Row(): | |
| # with gr.Column(scale=1, min_width=300): | |
| # input_files = gr.Files( | |
| # label="Upload Bracketed LDR Images", | |
| # file_types=["image"] | |
| # ) | |
| # process_button = gr.Button("Process Images", variant="primary") | |
| # with gr.Accordion("See Your Uploads", open=False): | |
| # input_gallery = gr.Gallery(label="Uploaded LDR Bracket", show_label=False, columns=3, height="auto") | |
| # with gr.Column(scale=2): | |
| # gr.Markdown("## Results") | |
| # with gr.Row(): | |
| # chosen_ldr_display = gr.Image(label="Best LDR Chosen by Algorithm", type="pil", interactive=False) | |
| # with gr.Row(): | |
| # output_cropped = gr.Image(label="Result 1: High-Quality Crop (728x728)", type="pil", interactive=False) | |
| # output_scaled = gr.Image(label="Result 2: Full Image (Scaled to 1024px Width)", type="pil", interactive=False) | |
| # process_button.click( | |
| # fn=process_images_and_display, | |
| # inputs=input_files, | |
| # outputs=[input_gallery, chosen_ldr_display, output_cropped, output_scaled] | |
| # ) | |
| # print("--- Launching Gradio App ---") | |
| # demo.launch(share=True) | |
| import gradio as gr | |
| import torch | |
| import os | |
| import tempfile | |
| import shutil | |
| from PIL import Image | |
| import numpy as np | |
| from pathlib import Path | |
| import sys | |
| import copy | |
| # --- Import logic from your project --- | |
| from options.test_options import TestOptions | |
| from data import create_dataset | |
| from models import create_model | |
| try: | |
| from best_ldr import compute_metrics_for_images, score_records | |
| except ImportError: | |
| raise ImportError("Could not import from best_ldr.py. Make sure the file is in the same directory as app.py.") | |
| print("--- Initializing LDR-to-HDR Model (this may take a moment) ---") | |
| # --- Global Setup: Load the CycleGAN model once when the app starts --- | |
| # We need to satisfy the parser's requirement for a dataroot at startup | |
| if '--dataroot' not in sys.argv: | |
| sys.argv.extend(['--dataroot', './dummy_dataroot_for_init']) | |
| # Load the base options | |
| opt = TestOptions().parse() | |
| # Manually override settings for our model | |
| opt.name = 'ldr2hdr_cyclegan_728' | |
| opt.model = 'test' | |
| opt.netG = 'resnet_9blocks' | |
| opt.norm = 'instance' | |
| opt.no_dropout = True | |
| opt.checkpoints_dir = './checkpoints' | |
| opt.gpu_ids = [0] if torch.cuda.is_available() else [] | |
| opt.device = torch.device('cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids else torch.device('cpu') | |
| # Create the model using these options | |
| model = create_model(opt) | |
| model.setup(opt) | |
| model.eval() | |
| print("--- Model Loaded Successfully ---") | |
| # --- The Main Gradio Processing Function --- | |
| def process_images_to_hdr(list_of_temp_files): | |
| """ | |
| The main workflow: select best LDR, run inference, and return results for the UI. | |
| """ | |
| if not list_of_temp_files: | |
| raise gr.Error("Please upload your bracketed LDR images.") | |
| if len(list_of_temp_files) < 2: | |
| gr.Warning("For best results, upload at least 2 bracketed LDR images.") | |
| uploaded_filepaths = [Path(f.name) for f in list_of_temp_files] | |
| try: | |
| # --- Step 1: Select the Best LDR --- | |
| print(f"Analyzing {len(uploaded_filepaths)} uploaded images...") | |
| weights = {"clipped": 0.35, "coverage": 0.25, "exposure": 0.15, "sharpness": 0.15, "noise": 0.10} | |
| records = compute_metrics_for_images(uploaded_filepaths, resize_max=1024) | |
| scored_records = score_records(records, weights) | |
| if not scored_records: | |
| raise gr.Error("Could not read or score any of the uploaded images.") | |
| best_ldr_record = scored_records[0] | |
| best_ldr_path = best_ldr_record['path'] | |
| print(f"Best LDR selected: {os.path.basename(best_ldr_path)} (Score: {best_ldr_record['score']:.4f})") | |
| # --- Step 2: Run Inference --- | |
| print("Running Full Image (High-Res Scaled) Inference...") | |
| # We only need the one set of options now | |
| inference_options = { | |
| 'preprocess': 'scale_width', | |
| 'load_size': 1024, # Generate the high-resolution, full image | |
| 'crop_size': 728 # This value is ignored but required by the parser | |
| } | |
| # Deep copy the base options to avoid modifying the global state | |
| local_opt = copy.deepcopy(opt) | |
| for key, value in inference_options.items(): | |
| setattr(local_opt, key, value) | |
| # Run the model | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| shutil.copy(best_ldr_path, temp_dir) | |
| local_opt.dataroot = temp_dir | |
| local_opt.num_test = 1 | |
| dataset = create_dataset(local_opt) | |
| for i, data in enumerate(dataset): | |
| model.set_input(data) | |
| model.test() | |
| visuals = model.get_current_visuals() | |
| for label, image_tensor in visuals.items(): | |
| if label == 'fake': | |
| image_numpy = (np.transpose(image_tensor.cpu().float().numpy()[0], (1, 2, 0)) + 1) / 2.0 * 255.0 | |
| final_hdr_image = Image.fromarray(image_numpy.astype(np.uint8)) | |
| print("Conversion to HDR successful.") | |
| # Return the gallery of inputs and the single final HDR image | |
| return uploaded_filepaths, final_hdr_image | |
| except Exception as e: | |
| print(f"An error occurred: {e}") | |
| raise gr.Error(f"An error occurred during processing: {e}") | |
| # --- Create and Launch the Gradio Interface --- | |
| with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}") as demo: | |
| gr.Markdown( | |
| """ | |
| # LDR Bracketing to HDR Converter | |
| Upload a set of bracketed LDR images. The app will automatically select the best one and convert it to a vibrant, full-resolution HDR image. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1, min_width=350): | |
| # --- INPUT --- | |
| input_files = gr.Files( | |
| label="Upload Bracketed LDR Images", | |
| file_types=["image"] | |
| ) | |
| process_button = gr.Button("Process Images", variant="primary") | |
| with gr.Accordion("See Your Uploaded Images", open=True): | |
| input_gallery = gr.Gallery(label="Uploaded Images", show_label=False, columns=[2, 3], height="auto") | |
| with gr.Column(scale=2): | |
| # --- OUTPUT --- | |
| gr.Markdown("## Generated HDR Result") | |
| output_image = gr.Image(label="Final HDR Image", type="pil", interactive=False, show_download_button=True) | |
| process_button.click( | |
| fn=process_images_to_hdr, | |
| inputs=input_files, | |
| outputs=[input_gallery, output_image] | |
| ) | |
| # gr.Markdown("### Examples") | |
| # gr.Examples( | |
| # examples=[ | |
| # [ | |
| # "../pix2pix_dataset/testA/077A2406.jpg", | |
| # "../pix2pix_dataset/testA/077A4049.jpg", | |
| # "../pix2pix_dataset/testA/077A4073.jpg" | |
| # ] | |
| # ], | |
| # inputs=input_files | |
| # ) | |
| print("--- Launching Gradio App ---") | |
| demo.launch(share=True) |