PyTorch绘制光流flow

最新推荐文章于 2024-06-07 09:49:54 发布

Evan_Tech

最新推荐文章于 2024-06-07 09:49:54 发布

阅读量1.9k

点赞数

CC 4.0 BY-SA版权

分类专栏： Python

本文链接：https://blog.youkuaiyun.com/yfren1123/article/details/104215553

本文介绍了两种在PyTorch中绘制光流flow的方法，重点是利用HSV色彩空间来表示flow的角度和长度。通过对比两种方法的输出结果，验证了第一种方法的正确性。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

flow
这里给出两种方法的code，第一种是用hsv来转化，第二种是直接绘制。

比较tricky的地方可能是思路上先生成hsv图，因为hsv图中第一个h维度，刚好用各种不同的颜色来表达角度，这个和flow中的角度是一致的。而flow的绝对值长度刚好又能被s和v维度表达（其实h和v只有一个自由度）

def arctan(x, y):
    y_x = abs(y/x)
    angle = (x >= 0) * (y > 0) * torch.atan(y_x) # first quadrant
    angle += (x < 0) * (y >= 0) * (torch.atan(y_x) + np.pi/2) # second quadrant
    angle += (x <= 0) * (y < 0) * (torch.atan(y_x) + np.pi) # third quadrant
    angle += (x > 0) * (y <= 0) * (torch.atan(y_x) + np.pi*3/2) # forth quadrant
    return angle


def flow_to_color(flow):

    mag = torch.sqrt(torch.pow(flow, 2).sum(dim=1)) # get abs flow
    max_mag_each_batch = torch.max(mag.reshape([mag.shape[0],-1]),1)[0]

    
    flow_u, flow_v = flow[:,0,:,:], flow[:,1,:,:]
    for i in range(len(max_mag_each_batch)):
        flow_u[i] /= max_mag_each_batch[i]
        flow_v[i] /= max_mag_each_batch[i]

    angle = arctan(flow_u, flow_v)

    im_h = angle / (2*np.pi)
    im_s = torch.clamp(mag / mag.max(), min=0, max=1) 
    im_v = torch.clamp(mag.max() - im_s, min=0