# 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)