batch_dehaze_alternatives.py

import os
import torch
import cv2
import numpy as np
from tqdm import tqdm
from PIL import Image
import torchvision.transforms as transforms
import shutil
import concurrent.futures

# Import Models
from src.dehazing.DehazeNet import DehazeNet
from src.dehazing.DCP import DCPDehazer

def get_dark_channel(img, patch_size=15):
    # img: [1, 3, H, W] tensor
    min_channel, _ = torch.min(img, dim=1, keepdim=True)
    # MinPool via negative MaxPool
    # patch_size should be odd
    pad = patch_size // 2
    dark_channel = -torch.nn.functional.max_pool2d(-min_channel, kernel_size=patch_size, stride=1, padding=pad)
    return dark_channel

def estimate_atmospheric_light(img, dark_channel):
    # img: [1, 3, H, W]
    # dark_channel: [1, 1, H, W]
    b, c, h, w = img.shape
    num_pixels = h * w
    num_search = int(max(num_pixels * 0.001, 1))
    
    dark_vec = dark_channel.view(b, num_pixels)
    img_vec = img.view(b, c, num_pixels)
    
    # Get top 0.1% indices
    _, indices = torch.topk(dark_vec, k=num_search, dim=1)
    
    # We want the pixel with highest intensity among these candidates
    # Since b=1 usually
    if b == 1:
        indices = indices[0]
        candidates = img_vec[0, :, indices] # [3, num_search]
        intensities = torch.sum(candidates, dim=0) # [num_search]
        max_idx = torch.argmax(intensities)
        A = candidates[:, max_idx].view(1, 3, 1, 1)
        return A
    else:
        # Batch processing (simplified: use mean of top 0.1%)
        # But for dehazing we usually process 1 by 1 or implement batch logic carefully
        # Here we just implement for b=1 as loop uses batch=1
        return torch.ones(b, 3, 1, 1).to(img.device)

def recover_image(img, transmission, A, t0=0.1):
    t = torch.clamp(transmission, min=t0)
    J = (img - A) / t + A
    return torch.clamp(J, 0, 1)

def process_dehazenet(input_dir, output_dir, weights_path, device='cuda'):
    print(f"\n--- DehazeNet Processing ---")
    device = torch.device(device if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    # Load Model
    model = DehazeNet().to(device)
    if os.path.exists(weights_path):
        model.load_state_dict(torch.load(weights_path, map_location=device))
        print(f"Loaded DehazeNet: {weights_path}")
    else:
        print(f"❌ Error: Weights not found at {weights_path}")
        return

    model.eval()
    
    os.makedirs(output_dir, exist_ok=True)
    
    # Prepare Labels Directory
    base_dir = os.path.dirname(input_dir)
    labels_src_dir = os.path.join(base_dir, 'labels')
    output_base_dir = os.path.dirname(output_dir)
    labels_dst_dir = os.path.join(output_base_dir, 'labels')
    os.makedirs(labels_dst_dir, exist_ok=True)
    
    files = [f for f in os.listdir(input_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp'))]
    transform = transforms.Compose([transforms.ToTensor()])
    
    with torch.no_grad():
        for filename in tqdm(files, desc="DehazeNet"):
            img_path = os.path.join(input_dir, filename)
            save_path = os.path.join(output_dir, filename)
            
            try:
                # DehazeNet often expects specific normalization, but basic ToTensor [0,1] is usually fine for inference
                pil_img = Image.open(img_path).convert('RGB')
                img_tensor = transform(pil_img).unsqueeze(0).to(device)
                
                # Model outputs Transmission Map
                transmission = model(img_tensor)
                
                # Estimate Atmospheric Light
                dark_channel = get_dark_channel(img_tensor)
                A = estimate_atmospheric_light(img_tensor, dark_channel)
                
                # Recover Image
                clean_tensor = recover_image(img_tensor, transmission, A)
                clean_img = clean_tensor.squeeze(0).cpu()
                
                clean_pil = transforms.ToPILImage()(clean_img)
                clean_pil.save(save_path)
                
                # Copy Label
                label_name = os.path.splitext(filename)[0] + '.txt'
                src_label = os.path.join(labels_src_dir, label_name)
                dst_label = os.path.join(labels_dst_dir, label_name)
                if os.path.exists(src_label):
                    shutil.copy2(src_label, dst_label)
                    
            except Exception as e:
                print(f"Error {filename}: {e}")

def process_dcp_single(args):
    filename, input_dir, output_dir, labels_src_dir, labels_dst_dir = args
    img_path = os.path.join(input_dir, filename)
    save_path = os.path.join(output_dir, filename)
    
    try:
        # Read Image (OpenCV uses BGR)
        # Handle Chinese characters in path by using numpy
        img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        if img is None: 
            print(f"Failed to read image: {img_path}")
            return
        
        dehazer = DCPDehazer() # Default params
        clean_img = dehazer.dehaze(img)
        
        # Save image (handle Chinese path)
        is_success, im_buf = cv2.imencode(os.path.splitext(filename)[1], clean_img)
        if is_success:
            im_buf.tofile(save_path)
        
        # Copy Label
        label_name = os.path.splitext(filename)[0] + '.txt'
        src_label = os.path.join(labels_src_dir, label_name)
        dst_label = os.path.join(labels_dst_dir, label_name)
        
        if os.path.exists(src_label):
            shutil.copy2(src_label, dst_label)
        else:
            # Try to find label if case sensitivity is issue or if it's missing
            pass
            
    except Exception as e:
        print(f"Error DCP {filename}: {e}")

def process_dcp(input_dir, output_dir):
    print(f"\n--- DCP Processing ---")
    os.makedirs(output_dir, exist_ok=True)
    
    base_dir = os.path.dirname(input_dir)
    labels_src_dir = os.path.join(base_dir, 'labels')
    output_base_dir = os.path.dirname(output_dir)
    labels_dst_dir = os.path.join(output_base_dir, 'labels')
    os.makedirs(labels_dst_dir, exist_ok=True)
    
    files = [f for f in os.listdir(input_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp'))]
    
    # Prepare args for multiprocessing
    tasks = [(f, input_dir, output_dir, labels_src_dir, labels_dst_dir) for f in files]
    
    # Use ProcessPoolExecutor for CPU-bound tasks
    # Limit workers to avoid memory overflow if images are large
    max_workers = min(os.cpu_count(), 8)
    print(f"Starting DCP with {max_workers} workers...")
    
    with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
        list(tqdm(executor.map(process_dcp_single, tasks), total=len(tasks), desc="DCP"))

if __name__ == "__main__":
    # Directories
    input_dir = r"data/RTTS/JPEGImages"
    
    # 1. DehazeNet
    dehazenet_out = r"data/RTTS_Dehazed_DehazeNet/images"
    dehazenet_weights = r"weights/DehazeNet_final.pth"
    process_dehazenet(input_dir, dehazenet_out, dehazenet_weights)
    
    # 2. DCP (Already done, skipping)
    # dcp_out = r"data/RTTS_Dehazed_DCP/images"
    # process_dcp(input_dir, dcp_out)