Initial commit of AHDRNet Gradio app

.gitattributes

@@ -0,0 +1 @@
+*.pth filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,141 @@
+import gradio as gr
+import torch
+import cv2
+import numpy as np
+import os
+from pathlib import Path
+from torchvision.transforms.functional import to_pil_image
+import argparse # Needed for the model initialization
+
+# --- 1. Import Your Model Architecture ---
+# This assumes your model.py file is in the same directory.
+from model import AHDR
+
+# --- 2. Define Constants and Helper Functions ---
+GAMMA = 2.2
+
+def load_and_preprocess_images(file_paths):
+    """
+    Loads images from temporary Gradio file paths, sorts them by brightness,
+    selects the 3 most representative ones, and preprocesses them for the model.
+    """
+    if len(file_paths) < 3:
+        raise gr.Error("Please upload at least 3 images for the best results. The model uses under, normal, and over-exposed shots.")
+
+    images_with_brightness = []
+    for file_path in file_paths:
+        img_bgr = cv2.imread(file_path)
+        if img_bgr is None:
+            continue
+        
+        # FIX 1: Corrected the OpenCV constant from COLOR_BGR2GRAYSCALE to COLOR_BGR2GRAY
+        gray_img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
+        brightness = np.mean(gray_img)
+        images_with_brightness.append({'path': file_path, 'brightness': brightness, 'img': img_bgr})
+    
+    if len(images_with_brightness) < 3:
+        raise gr.Error(f"Could only read {len(images_with_brightness)} valid images, but need at least 3.")
+
+    images_with_brightness.sort(key=lambda x: x['brightness'])
+    
+    selected_images = [
+        images_with_brightness[0]['img'],
+        images_with_brightness[len(images_with_brightness) // 2]['img'],
+        images_with_brightness[-1]['img']
+    ]
+    
+    proxy_exposure_times = [0.25, 1.0, 4.0] 
+    model_inputs = []
+    for i, ldr_bgr in enumerate(selected_images):
+        ldr_rgb = cv2.cvtColor(ldr_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
+        exposure = proxy_exposure_times[i]
+        hdr_from_ldr = np.power(ldr_rgb, GAMMA) / exposure
+        six_ch_input = np.concatenate((ldr_rgb, hdr_from_ldr), axis=2)
+        tensor_input = torch.from_numpy(six_ch_input.transpose(2, 0, 1)).float()
+        model_inputs.append(tensor_input)
+        
+    return model_inputs[0].unsqueeze(0), model_inputs[1].unsqueeze(0), model_inputs[2].unsqueeze(0)
+
+def postprocess_output(tensor):
+    """Converts the model's output tensor to a displayable PIL image."""
+    tensor = torch.clamp(tensor, 0, 1)
+    return to_pil_image(tensor.squeeze(0).cpu())
+
+# --- 3. Load The Model (Done Once at Startup) ---
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+CHECKPOINT_PATH = "./finetuning_run_01/checkpoints/best_model.pth"
+
+model_args = argparse.Namespace(nDenselayer=6, growthRate=32, nBlock=16, nFeat=64, nChannel=6)
+model = AHDR(model_args).to(DEVICE)
+
+print(f"Loading checkpoint from: {CHECKPOINT_PATH}")
+checkpoint = torch.load(CHECKPOINT_PATH, map_location=DEVICE)
+
+state_dict = checkpoint.get('state_dict', checkpoint)
+state_dict_cleaned = {k.replace('module.', ''): v for k, v in state_dict.items()}
+model.load_state_dict(state_dict_cleaned)
+model.eval()
+print("Model loaded successfully and is in evaluation mode.")
+
+# --- 4. Define the Gradio Inference Function ---
+def generate_hdr(file_paths): # FIX 2: The input is now a list of file paths directly
+    """
+    The main function that Gradio will call. It takes a list of uploaded file paths,
+    processes them, and returns the HDR image.
+    """
+    if file_paths is None or len(file_paths) < 3:
+         raise gr.Error("Please upload at least 3 images.")
+         
+    print(f"Processing {len(file_paths)} uploaded images...")
+    
+    try:
+        i1, i2, i3 = load_and_preprocess_images(file_paths) # Pass the list of paths directly
+        i1, i2, i3 = i1.to(DEVICE), i2.to(DEVICE), i3.to(DEVICE)
+
+        with torch.no_grad():
+            prediction_tensor = model(i1, i2, i3)
+        
+        output_image = postprocess_output(prediction_tensor)
+        
+        print("Successfully generated HDR image.")
+        return output_image
+        
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        raise gr.Error(f"An error occurred during processing: {e}")
+
+# --- 5. Create and Launch the Gradio Interface ---
+title = "Attention-guided HDR (AHDRNet)"
+description = """
+Upload a set of multi-exposure LDR images (at least 3 are recommended: under-exposed, normally-exposed, and over-exposed) of the same scene. 
+The model will automatically select the three most representative images and merge them into a single, well-exposed High Dynamic Range (HDR) image.
+"""
+
+example_path = Path("examples")
+examples = [
+    [str(p) for p in example_path.glob(f"{d.name}/*")]
+    for d in example_path.iterdir()
+    if d.is_dir()
+] if example_path.exists() else []
+
+iface = gr.Interface(
+    fn=generate_hdr,
+    inputs=gr.File(
+        file_count="multiple",
+        label="Upload LDR Images (JPG, PNG)",
+        type="filepath" 
+    ),
+    outputs=gr.Image(type="pil", label="Generated HDR Image"),
+    title=title,
+    description=description,
+    examples=examples,
+    # FIX 3: Removed deprecated 'allow_flagging' parameter
+)
+
+# if __name__ == "__main__":
+#     # FIX 4: Use share=True for temporary local testing. 
+#     # REMOVE share=True when uploading to Hugging Face Spaces.
+#     iface.launch(share=True)
+
+if __name__ == "__main__":
+    iface.launch()

finetuning_run_01/checkpoints/best_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1a1c03ba47c833deeafc9bc52a46780a1a35060653d864a341ae8a7712298eb
+size 17668173

model.py

@@ -0,0 +1,105 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F 
+import numpy as np
+from torch.autograd import Variable
+
+        
+class make_dilation_dense(nn.Module):
+  def __init__(self, nChannels, growthRate, kernel_size=3):
+    super(make_dilation_dense, self).__init__()
+    self.conv = nn.Conv2d(nChannels, growthRate, kernel_size=kernel_size, padding=(kernel_size-1)//2+1, bias=True, dilation=2)
+  def forward(self, x):
+    out = F.relu(self.conv(x))
+    out = torch.cat((x, out), 1)
+    return out
+
+# Dilation Residual dense block (DRDB)
+class DRDB(nn.Module):
+  def __init__(self, nChannels, nDenselayer, growthRate):
+    super(DRDB, self).__init__()
+    nChannels_ = nChannels
+    modules = []
+    for i in range(nDenselayer):    
+        modules.append(make_dilation_dense(nChannels_, growthRate))
+        nChannels_ += growthRate 
+    self.dense_layers = nn.Sequential(*modules)    
+    self.conv_1x1 = nn.Conv2d(nChannels_, nChannels, kernel_size=1, padding=0, bias=True)
+  def forward(self, x):
+    out = self.dense_layers(x)
+    out = self.conv_1x1(out)
+    out = out + x
+    return out
+
+# Attention Guided HDR, AHDR-Net
+class AHDR(nn.Module):
+    def __init__(self, args):
+        super(AHDR, self).__init__()
+        nChannel = args.nChannel
+        nDenselayer = args.nDenselayer
+        nFeat = args.nFeat
+        growthRate = args.growthRate
+        self.args = args
+
+        # F-1
+        self.conv1 = nn.Conv2d(nChannel, nFeat, kernel_size=3, padding=1, bias=True)
+        # F0
+        self.conv2 = nn.Conv2d(nFeat*3, nFeat, kernel_size=3, padding=1, bias=True)
+        self.att11 = nn.Conv2d(nFeat*2, nFeat*2, kernel_size=3, padding=1, bias=True)
+        self.att12 = nn.Conv2d(nFeat*2, nFeat, kernel_size=3, padding=1, bias=True)
+        self.attConv1 = nn.Conv2d(nFeat, nFeat, kernel_size=3, padding=1, bias=True)
+        self.att31 = nn.Conv2d(nFeat*2, nFeat*2, kernel_size=3, padding=1, bias=True)
+        self.att32 = nn.Conv2d(nFeat*2, nFeat, kernel_size=3, padding=1, bias=True)
+        self.attConv3 = nn.Conv2d(nFeat, nFeat, kernel_size=3, padding=1, bias=True)
+
+        # DRDBs 3
+        self.RDB1 = DRDB(nFeat, nDenselayer, growthRate)
+        self.RDB2 = DRDB(nFeat, nDenselayer, growthRate)
+
+        self.RDB3 = DRDB(nFeat, nDenselayer, growthRate)
+        # feature fusion (GFF)
+        self.GFF_1x1 = nn.Conv2d(nFeat*3, nFeat, kernel_size=1, padding=0, bias=True)
+        self.GFF_3x3 = nn.Conv2d(nFeat, nFeat, kernel_size=3, padding=1, bias=True)
+        # fusion
+        self.conv_up = nn.Conv2d(nFeat, nFeat, kernel_size=3, padding=1, bias=True)
+
+        # conv 
+        self.conv3 = nn.Conv2d(nFeat, 3, kernel_size=3, padding=1, bias=True)
+        self.relu = nn.LeakyReLU()
+
+
+    def forward(self, x1, x2, x3):
+
+        F1_ = self.relu(self.conv1(x1))
+        F2_ = self.relu(self.conv1(x2))
+        F3_ = self.relu(self.conv1(x3))
+
+        F1_i = torch.cat((F1_, F2_), 1)
+        F1_A = self.relu(self.att11(F1_i))
+        F1_A = self.att12(F1_A)
+        F1_A = nn.functional.sigmoid(F1_A)
+        F1_ = F1_ * F1_A
+
+
+        F3_i = torch.cat((F3_, F2_), 1)
+        F3_A = self.relu(self.att31(F3_i))
+        F3_A = self.att32(F3_A)
+        F3_A = nn.functional.sigmoid(F3_A)
+        F3_ = F3_ * F3_A
+
+        F_ = torch.cat((F1_, F2_, F3_), 1)
+
+        F_0 = self.conv2(F_)
+        F_1 = self.RDB1(F_0)
+        F_2 = self.RDB2(F_1)
+        F_3 = self.RDB3(F_2)
+        FF = torch.cat((F_1, F_2, F_3), 1)
+        FdLF = self.GFF_1x1(FF)         
+        FGF = self.GFF_3x3(FdLF)
+        FDF = FGF + F2_
+        us = self.conv_up(FDF)
+
+        output = self.conv3(us)
+        output = nn.functional.sigmoid(output)
+
+        return output

requirements.txt

@@ -0,0 +1,5 @@
+torch
+torchvision
+gradio
+opencv-python-headless
+numpy

utils.py

@@ -0,0 +1,12 @@
+import os
+import torch
+
+def mk_dir(dir_path):
+    if not os.path.exists(dir_path):
+        os.makedirs(dir_path)
+
+def model_load(model, trained_model_dir, model_file_name):
+    model_path = os.path.join(trained_model_dir, model_file_name)
+    # trained_model_dir + model_file_name    # '/modelParas.pkl'
+    model.load_state_dict(torch.load(model_path))
+    return model