[FIXED] overlaying the ground truth mask on an image

Issue

In my project, I extracted frames from a video and in another folder I have ground truth for each frame. I want to map the ground truth image of each frame of a video (in my case, it is saliency prediction ground truth) on its related frame image. As an example I have the following frame:

And the following is ground truth mask:

and the following is the mapping of ground truth on the frame.

How can I do that. Also, I have two folders that inside each of them, there are several folders that inside each of them the there are stored frames. How can I do this operation with these batch data?

This is the hierarchy of my folders:

frame_folder: folder_1, folder_2, ......

├── frames
│   ├── 601   (601 and 602  and etc are folders that in the inside there are image frames that their name is like 0001.png,0002.png, ...)
│   ├── 602
       .
       .
       .
│   └── 700


 ├── ground truth
    │   ├── 601   (601 and 602  and etc are folders that in the inside there are ground truth masks that their name is like 0001.png,0002.png, ...)
    │   ├── 602
           .
           .
           .
    │   └── 700

Update: Using the answer proposed by @hkchengrex , I faced with an error. When there is only one folder in the paths, it works well but when I put several folders (frames of different videos) based on the question I face with the following error. the details are in below:

 multiprocessing.pool.RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/home/user/miniconda3/envs/vtn/lib/python3.10/multiprocessing/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
TypeError: process_video() takes 1 positional argument but 6 were given
"""

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/user/Video_processing/Saliency_mapping.py", line 69, in <module>
    pool.apply(process_video, videos)
  File "/home/user/miniconda3/envs/vtn/lib/python3.10/multiprocessing/pool.py", line 357, in apply
    return self.apply_async(func, args, kwds).get()
  File "/home/user/miniconda3/envs/vtn/lib/python3.10/multiprocessing/pool.py", line 771, in get
    raise self._value
TypeError: process_video() takes 1 positional argument but 6 were given

Solution

I need to do similar things pretty often. In my favorite StackOverflow fashion, here is a script that you can copy and paste. I hope the code itself is self-explanatory. There are a few things that you can tune and try (e.g., color maps, overlay styles). It uses multiprocessing.Pool for faster batch-processing, resizes the mask to match the shape of the image, assumes the mask is in .png format, and depends on the file structure that you posted.

import os
from os import path
import cv2
import numpy as np

from argparse import ArgumentParser
from multiprocessing import Pool


def create_overlay(image, mask):
    """
    image: H*W*3 numpy array
    mask: H*W numpy array
    If dimensions do not match, the mask is upsampled to match that of the image

    Returns a H*W*3 numpy array
    """
    h, w = image.shape[:2]
    mask = cv2.resize(mask, dsize=(w,h), interpolation=cv2.INTER_CUBIC)

    # color options: https://docs.opencv.org/4.x/d3/d50/group__imgproc__colormap.html
    mask_color = cv2.applyColorMap(mask, cv2.COLORMAP_HOT).astype(np.float32)
    mask = mask[:, :, None] # create trailing dimension for broadcasting
    mask = mask.astype(np.float32)/255

    # different other options that you can use to merge image/mask
    overlay = (image*(1-mask)+mask_color*mask).astype(np.uint8)
    # overlay = (image*0.5 + mask_color*0.5).astype(np.uint8)
    # overlay = (image + mask_color).clip(0,255).astype(np.uint8)

    return overlay

def process_video(video_name):
    """
    Processing frames in a single video
    """
    vid_image_path = path.join(image_path, video_name)
    vid_mask_path = path.join(mask_path, video_name)
    vid_output_path = path.join(output_path, video_name)
    os.makedirs(vid_output_path, exist_ok=True)

    frames = sorted(os.listdir(vid_image_path))
    for f in frames:
        image = cv2.imread(path.join(vid_image_path, f))
        mask = cv2.imread(path.join(vid_mask_path, f.replace('.jpg','.png')), cv2.IMREAD_GRAYSCALE)
        overlay = create_overlay(image, mask)
        cv2.imwrite(path.join(vid_output_path, f), overlay)


parser = ArgumentParser()
parser.add_argument('--image_path')
parser.add_argument('--mask_path')
parser.add_argument('--output_path')
args = parser.parse_args()

image_path = args.image_path
mask_path = args.mask_path
output_path = args.output_path

if __name__ == '__main__':
    videos = sorted(
        list(set(os.listdir(image_path)).intersection(
                set(os.listdir(mask_path))))
    )

    print(f'Processing {len(videos)} videos.')

    pool = Pool()
    pool.map(process_video, videos)

    print('Done.')

Output:

EDIT: Made it work on Windows; changed pool.apply to pool.map.

Answered By - hkchengrex