我是Python新手，目前正处于图像处理的学习阶段。由于我需要处理大型数据集，而这些数据集基本上只需要阈值处理，因此通过使用GPU，TensorFlow或PyTorch可能更适合。有相关示例吗？我应该怎么做？
下面是我目前使用OpenCV的方法。

import numpy as np
import cv2
import os

folderDir = "C://Users/ruler/Desktop/testseg/"

total = []

with open('readme.txt', 'w') as f:
    count = 0
    for allImages in os.listdir(folderDir):

        if (allImages.startswith('TRAIN_SET') and allImages.endswith(".bmp")):

            img = cv2.imread(os.path.join(folderDir, allImages))

            gry = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

            _,blackMask = cv2.threshold(gry, 0, 255, cv2.THRESH_BINARY_INV)
            _, blackgreyMask = cv2.threshold(gry, 128, 255, cv2.THRESH_BINARY_INV)
            whiteMask = cv2.bitwise_not(blackgreyMask)
            greyMask = cv2.bitwise_xor(blackMask, blackgreyMask)

            x1 = []
            y1 = []

            def verticle(mask, y, x):
                vertiPixel = 0
                while(y < mask.shape[0]):
                    if (y + 1) == mask.shape[0]:
                        break
                    else:
                        if(mask[y + 1][x] == 255):
                            vertiPixel += 1
                            y += 1
                        else:
                            break

                y1.append(vertiPixel)
                return y1

            def horizontal(mask, y, x):
                horiPixel = 0
                while(x < mask.shape[1]):
                    if (x + 1) == mask.shape[1]:
                        break
                    else:
                        if(mask[y][x + 1] == 255):
                            horiPixel += 1
                            x += 1
                        else:
                            break

                x1.append(horiPixel)
                return x1

            def mask(mask):

                for y in range (mask.shape[0]):
                    for x in range (mask.shape[1]):

                        if(mask[y][x] == 255):
                            verticle(mask, y, x)
                            horizontal(mask, y, x)

mask(blackMask)
print(np.average(y1), np.average(x1))

我尝试查找与阈值设置相关的TensorFlow或PyTorch，但没有找到。

期望

通过TensorFlow或PyTorch实现了阈值分割API，加快了图像阈值分割的速度。