RAFT Optical Flow With Greyscale Images

Chance_Lawrence: <p>Hi,</p>
<p>I am new to this form, so I am not sure all the info that is expected from me in in this post.</p>
<p>I am currently collecting data with an Allied Vision Alvium 1800 U-511m Monochrome Camera. Earlier in the week, I had successfully demonstrated optical flow visualization with RAFT using color imagery on other cameras, but I have been unable to recreate that same luck with monochrome data.</p>
<p>I am using torchvision currently, and I have tried to edit the dimensions of my image tensors to fit with the expected 3 color channels for RAFT, but my output flow is a jumbled, incoherent mess.</p>
<p>I start by reading in 4 images (The camera returns the data as individual monochrome jpgs, so I am selecting the frame manually in a folder)</p>
<p>frame_indx is just a list containing a number of arbitrary frame index pairs that I want to calculate flow over</p>
<pre data-code-wrap="python"><code class="lang-python">
frame_indx = [40,41,50,51]

frames = torch.empty(size=(4,3,2064,2464))

for i in range(len(frame_indx)):
    frame = read_image(f'tmp/fish_eye_test/VmbImage_{frame_indx[i]}.jpg')
    frame = torch.Tensor.expand(frame,[3,2064,2464])
    frames[i] = frame
</code></pre>
<p>I am forcing the 1 channel monochrome into 3 color channels. Does this mess with the raft_large model? If the model has been trained on color data, is it even possible to implement monochrome without providing my own data set? Are there optical flow algorithms for monochrome data?</p>
<p>Here is the rest of my code for completeness</p>
<pre data-code-wrap="python"><code class="lang-python">import numpy as np
import torch
import matplotlib.pyplot as plt
import torchvision.transforms.functional as F
import torchvision.transforms as T
    
def plot(imgs, **imshow_kwargs):
    if not isinstance(imgs[0], list):
        # Make a 2d grid even if there's just 1 row
        imgs = [imgs]

    num_rows = len(imgs)
    num_cols = len(imgs[0])
    _, axs = plt.subplots(nrows=num_rows, ncols=num_cols, squeeze=False)
    for row_idx, row in enumerate(imgs):
        for col_idx, img in enumerate(row):
            ax = axs[row_idx, col_idx]
            ax.imshow(np.asarray(img)[0,:,:],cmap='grey', **imshow_kwargs)
            ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])

    plt.tight_layout()
    
    
from torchvision.io import read_image

frame_indx = [40,41,50,51]
frames = torch.empty(size=(4,3,2064,2464))

for i in range(len(frame_indx)):
    frame = read_image(f'tmp/fish_eye_test/VmbImage_{frame_indx[i]}.jpg')
    frame = torch.Tensor.expand(frame,[3,2064,2464])
    frames[i] = frame


img1_batch = torch.stack([frames[0], frames[2]])
img2_batch = torch.stack([frames[1], frames[3]])

def preprocess(batch):
    transforms = T.Compose(
        [
            T.ConvertImageDtype(torch.float32),
            T.Normalize(mean=0.5, std=0.5),  # map [0, 1] into [-1, 1]
            T.Resize(size=(240, 320)),
        ]
    )
    batch = transforms(batch)
    return batch

device = "cuda" if torch.cuda.is_available() else "cpu"

img1_batch = preprocess(img1_batch).to(device)
img2_batch = preprocess(img2_batch).to(device)

from torchvision.models.optical_flow import raft_large

model = raft_large(pretrained=True, progress=False).to(device)
model = model.eval()

list_of_flows = model(img1_batch.to(device), img2_batch.to(device))
predicted_flows = list_of_flows[-1]

from torchvision.utils import flow_to_image

flow_imgs = flow_to_image(predicted_flows)


# The images have been mapped into [-1, 1] but for plotting we want them in [0, 1]
img1_batch = [(img1 + 1) / 2 for img1 in img1_batch]


pred_flow_grid = [[img1, pred_flow] for (img1, pred_flow) in zip(img1_batch, predicted_flows)]


grid = [[img1, flow_img] for (img1, flow_img) in zip(img1_batch, flow_imgs)]
plot(grid)
</code></pre>
<p>And the visualized output is as follows</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/1/d1c229cbd18291d18d33b0a31800d37ccefea31e.png" alt="Figure 2024-10-03 174334" data-base62-sha1="tVBMCtBt83DijTY0Q3DFbXa3DNY" width="385" height="280"></p>
<p>TLDR: Monochrome data is not working for RAFT optical flow. Not sure how to alter my data so that I can get it working or if it is even possible</p>