我已将其转换为使用Opencv库，而关于如何实时从摄像头读取图像处理的问题

[NowLoadY]

您好，这是我更改后的test.py代码：

import cv2
import numpy as np
import os
import torch
import torch.utils
from torch.autograd import Variable
from SCI_model import Finetunemodel


class MemoryFriendlyLoader(torch.utils.data.Dataset):
    def __init__(self, img_dir):
        self.low_img_dir = img_dir
        self.train_low_data_names = []

        for root, dirs, names in os.walk(self.low_img_dir):
            for name in names:
                self.train_low_data_names.append(os.path.join(root, name))

        self.train_low_data_names.sort()
        self.count = len(self.train_low_data_names)

    def transform(self, image):
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = image.transpose((2, 0, 1))
        image = image.astype(np.float32) / 255.0
        image = torch.from_numpy(image)
        return image

    def load_images_transform(self, file):
        # 使用 OpenCV 读入图片
        im_cv = cv2.imread(file)
        # 调整图像大小
        im_np = cv2.resize(im_cv, (600, 400))
        # 将 NumPy 数组转换为 PyTorch Tensor
        im = self.transform(im_np)
        return im

    def __getitem__(self, index):
        low_file = self.train_low_data_names[index]
        low = self.load_images_transform(low_file)
        return low, low_file

    def __len__(self):
        return self.count


def main():
    modelPath = './SCI_weights/medium.pt'
    savePath = './SCI_results/medium'
    dataPath = './SCI_data/medium'
    os.makedirs(savePath, exist_ok=True)
    TestDataset = MemoryFriendlyLoader(img_dir=dataPath)
    test_queue = torch.utils.data.DataLoader(
        TestDataset, batch_size=1,
        pin_memory=True, num_workers=0)
    model = Finetunemodel(modelPath)
    model = model.cuda()
    model.eval()
    with torch.no_grad():
        for _, (input, image_name) in enumerate(test_queue):
            input = Variable(input, volatile=True).cuda()
            i, r = model(input)
            image_numpy = r[0].cpu().float().numpy()
            image_numpy = (np.transpose(image_numpy, (1, 2, 0)))
            image_numpy = np.clip(image_numpy * 255, 0, 255)
            image_cv = cv2.cvtColor(image_numpy, cv2.COLOR_RGB2BGR).astype(np.uint8)  # 将 NumPy 数组转换为 OpenCV 图片
            cv2.imshow("result", image_cv)
            cv2.waitKey(0)


if __name__ == '__main__':
    main()

[NowLoadY]

尝试用下面的代码实现了实时增强，但是很明显不够规范：

import cv2
import numpy as np
import os
import torch
import torch.utils
from torch.autograd import Variable
from SCI_model import Finetunemodel


class MemoryFriendlyLoader(torch.utils.data.Dataset):
    def __init__(self, img_dir):
        self.low_img_dir = img_dir
        self.train_low_data_names = []

        for root, dirs, names in os.walk(self.low_img_dir):
            for name in names:
                self.train_low_data_names.append(os.path.join(root, name))

        self.train_low_data_names.sort()
        self.count = len(self.train_low_data_names)

    def transform(self, image):
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = image.transpose((2, 0, 1))
        image = image.astype(np.float32) / 255.0
        image = torch.from_numpy(image)
        return image

    def load_images_transform(self, file):
        # 使用 OpenCV 读入图片
        #im_cv = cv2.imread(file)
        _, im_cv = cam.read()
        cv2.imshow("ori", im_cv)
        # 调整图像大小
        #im_np = cv2.resize(im_cv, (600, 400))
        # 将 NumPy 数组转换为 PyTorch Tensor
        im = self.transform(im_cv)
        return im

    def __getitem__(self, index):
        low_file = self.train_low_data_names[index]
        low = self.load_images_transform(low_file)
        return low, low_file

    def __len__(self):
        return self.count


def main():
    modelPath = './SCI_weights/difficult.pt'
    savePath = './SCI_results/medium'
    dataPath = './SCI_data/medium'
    os.makedirs(savePath, exist_ok=True)
    TestDataset = MemoryFriendlyLoader(img_dir=dataPath)
    test_queue = torch.utils.data.DataLoader(
        TestDataset, batch_size=1,
        pin_memory=True, num_workers=0)
    model = Finetunemodel(modelPath)
    model = model.cuda()
    model.eval()

    while True:
        with torch.no_grad():
            for _, (input, image_name) in enumerate(test_queue):
                input = Variable(input, volatile=True).cuda()
                i, r = model(input)
                image_numpy = r[0].cpu().float().numpy()
                image_numpy = (np.transpose(image_numpy, (1, 2, 0)))
                image_numpy = np.clip(image_numpy * 255, 0, 255)
                image_cv = cv2.cvtColor(image_numpy, cv2.COLOR_RGB2BGR).astype(np.uint8)  # 将 NumPy 数组转换为 OpenCV 图片
                cv2.imshow("result", image_cv)
                cv2.waitKey(1)


if __name__ == '__main__':
    cam = cv2.VideoCapture(0)
    main()