nir_rgb_merge.py

from glob import glob
from PIL import Image
from tqdm import tqdm
import multiprocessing
import numpy as np

nir_folder = "/home/prateekjha/anomaly_detection/full_data/Agriculture-Vision-2021/val/images/nir"
rgb_folder = "/home/prateekjha/anomaly_detection/full_data/Agriculture-Vision-2021/val/images/rgb"
rgb_files = glob(f"/home/prateekjha/anomaly_detection/full_data/Agriculture-Vision-2021/val/images/rgb/*")
save_folder = "/home/prateekjha/anomaly_detection/full_data/Agriculture-Vision-2021/val/images/rgb_nir"
rgb_ids = [x.split("/")[-1] for x in rgb_files]

nir_files = [f"{nir_folder}/{idx}" for idx in rgb_ids]

# Define a function to process a single image
def process_image(idx):
    nir_image_path = nir_folder + "/" + idx
    rgb_image_path = rgb_folder + "/" + idx

    nir_image = np.expand_dims(np.array(Image.open(nir_image_path)), 2)
    rgb_image = np.array(Image.open(rgb_image_path))

    rgb_nir = np.concatenate([rgb_image, nir_image], 2)

    im = Image.fromarray(rgb_nir)
    idx = idx.replace("jpg", "png")
    im.save(f"{save_folder}/{idx}")

if __name__ == "__main__":
    # Create a multiprocessing pool
    num_processes = multiprocessing.cpu_count()
    pool = multiprocessing.Pool(processes=num_processes)

    # Use tqdm to add a progress bar
    with tqdm(total=len(rgb_ids)) as pbar:
        for _ in pool.imap_unordered(process_image, rgb_ids):
            pbar.update(1)

    pool.close()
    pool.join()