[FIXED] Merge images by mask

Issue

I'm trying to merge two images according to the values in a mask where at all points where the mask is 1, the resulting image has the values of the first image, and otherwise it has values of the second image. Does anyone know how it can be achieved in pytorch? Using numpy, it can be achieved using

>>> import numpy as np
>>> img1 = np.random.rand(100,100,3)
>>> img2 = np.random.rand(100,100,3)
>>> mask = np.random.rand(100,100)>.5
>>> res = img2.copy()
>>> res[mask] = img1[mask]

Solution

You are looking for np.where:

res = no.where(mask, img1, img2)

Answered By - Ivan