for (img1_file, img2_file) in tqdm(img_pairs): img1 = np.array(imread(img1_file)) img2 = np.array(imread(img2_file)) if args.arch == 'StrainNet_l' and img1.ndim == 3: img1 = img1[:,:,1] img2 = img2[:,:,1] img1 = img1/255 img2 = img2/255 if img1.ndim == 2: img1 = img1[np.newaxis, ...] img2 = img2[np.newaxis, ...] img1 = img1[np.newaxis, ...] img2 = img2[np.newaxis, ...] img1 = torch.from_numpy(img1).float() img2 = torch.from_numpy(img2).float() if args.arch == 'StrainNet_h' or args.arch == 'StrainNet_f': img1 = torch.cat([img1,img1,img1],1) img2 = torch.cat([img2,img2,img2],1) input_var = torch.cat([img1,img2],1) elif img1.ndim == 3: img1 = np.transpose(img1, (2, 0, 1)) img2 = np.transpose(img2, (2, 0, 1)) img1 = torch.from_numpy(img1).float() img2 = torch.from_numpy(img2).float() input_var = torch.cat([img1, img2]).unsqueeze(0) # compute output input_var = input_var.to(device) output = model(input_var) if args.arch == 'StrainNet_h' or args.arch == 'StrainNet_l': output = torch.nn.functional.interpolate(input=output, scale_factor=2, mode='bilinear') output_to_write = output.data.cpu() output_to_write = output_to_write.numpy() disp_x = output_to_write[0,0,:,:] disp_x = - disp_x * args.div_flow + 1 disp_y = output_to_write[0,1,:,:] disp_y = - disp_y * args.div_flow + 1 filenamex = save_path/'{}{}'.format(img1_file.stem[:-1], '_disp_x') filenamey = save_path/'{}{}'.format(img1_file.stem[:-1], '_disp_y') np.savetxt(filenamex + '.csv', disp_x,delimiter=',') np.savetxt(filenamey + '.csv', disp_y,delimiter=',')
时间: 2024-02-14 10:25:06 浏览: 163
PairsTrading_FEX.zip_Pairs_PairsTrading_FEX_pairs trading_tradin
这代码是一个图像处理的代码片段,它的作用是对一对图像进行处理并输出结果。
首先,代码使用imread函数读取两个图像文件(img1_file和img2_file),然后将其转换为numpy数组(img1和img2)。
接下来,根据参数args.arch的值和图像的维度,对图像进行一些预处理操作。如果args.arch等于'StrainNet_l'并且图像是三维的,则只保留第二个通道。然后将图像的像素值归一化到0到1之间。
然后,根据图像的维度进行不同的处理。如果图像是二维的,则添加一个额外的维度,并将其转换为Tensor类型。如果args.arch等于'StrainNet_h'或'StrainNet_f',则将图像在通道维度上进行复制。最后,将两个图像在通道维度上拼接起来,得到input_var。
如果图像是三维的,则对其进行转置操作,并转换为Tensor类型。然后将两个图像拼接起来,并在第0维度上添加一个额外的维度,得到input_var。
接下来,将input_var传入模型(model)进行计算得到输出(output)。如果args.arch等于'StrainNet_h'或'StrainNet_l',则对输出进行双线性插值操作。
然后,将输出转移到CPU上,并转换为numpy数组(output_to_write)。根据需要,将输出进行一些后处理操作,最终得到disp_x和disp_y。
最后,将disp_x和disp_y保存为CSV文件,文件名根据输入图像的文件名生成,并保存在save_path路径下。
以上就是这段代码的功能和流程。如果有其他问题,请随时提问!
阅读全文