How to get seams of a color image in OpenCV2 in Python?

Question

I am writing a bit of code to allow me to remove seams from an image. Currently, my code allows me to find, highlight, and remove seams from a greyscale image. I am needing to remove seams from colored images and my code will not allow me to do that. Can anyone help me modify my code to allow me to do this?

My code:

import numpy as np
import cv2
import math
import time

def getEdgeImage(img,margin=10):
    kernel=np.float64([[-1,0,1]])
    Ix=cv2.filter2D(img,cv2.CV_64F,kernel)
    Iy=cv2.filter2D(img,cv2.CV_64F,kernel)
    I=np.hypot(Ix,Iy)
    m=I.max()
    I[:,:margin]=m
    I[:,-margin:]=m
    return I

def getEnergyMap(img,repulseMask=None,attractMask=None):
    edges=getEdgeImage(img)
    if attractMask is not None:
        edges[attractMask==1]=-10
    if repulseMask is not None:
        edges[repulseMask==1]=235
    kernel=np.ones(3,np.float64)
    for i in range(1,len(edges)):
        minAbove=cv2.erode(edges[i-1],kernel).T[0]
        edges[i]+=minAbove
    return edges

def getSeam(img,repulseMask=None,attractMask=None):
    energyMap=getEnergyMap(img,repulseMask,attractMask)
    y=len(energyMap)-1
    x=np.argmin(energyMap[y])
    seam=[(x,y)]
    while len(seam)<len(energyMap):
        x,y=seam[-1]
        newY=y-1
        newX=x+np.argmin(energyMap[newY,x-1:x+2])-1
        seam.append((newX,newY))
    return seam

img1=cv2.imread("image.jpg") #[::2,::2]
# attractMask=img1*0
# repulseMask=img1*0
seam=getSeam(img1)

The attract and repulse masks are unimportant to the code currently, they are just used so I can manually plug in pixel coordinates to increase or decrease the amount of seams going in that coordinate plane.

The error I get when I run this code:

Traceback (most recent call last):
  File "Program.py", line 110, in <module>
    seam=getSeam(img1)
  File "Program.py", line 62, in getSeam
    energyMap=getEnergyMap(img,repulseMask,attractMask)
  File "Program.py", line 58, in getEnergyMap
    edges[i]+=minAbove
ValueError: operands could not be broadcast together with shapes (960,3) (960,) (960,3)

Is there anyway that I can get this to work with my code? I'll modify the functions if that is what I need to do.

Answer 1

Then try this, these are separate channels given to functions individually.

r=img1[:,:,0]
seam_r=getSeam(r)

g=img1[:,:,1]
seam_g=getSeam(g)

b=img1[:,:,2]
seam_b=getSeam(b)

After this, pass the results to your post function individually.