python 读图片封装 支持 视频，图片文件夹，图片

2024.02.01更新

安装依赖项：pip install natsort

#-*-coding:utf-8-*-
import os.path
from natsort import natsorted
import cv2class ImgReader:def __init__(self, source, suffix='mp4'):self.suffix = suffixif suffix!='mp4':self.type = 'file'if os.path.isdir(source):img_files = ['%s/%s' % (i[0].replace("\\", "/"), j) for i in os.walk(source) for j in i[-1] if j.endswith((type))]self.img_files = natsorted(img_files)self.img_index = 0self.type = 'dir'self.total_frames = len(self.img_files)elif source.endswith(".mp4"):self.type = 'mp4'self.cap = cv2.VideoCapture(source)if not self.cap.isOpened():raise ValueError(f"Error: Could not open video file at {source}")self.fps=int(self.cap.get(cv2.CAP_PROP_FPS)) + 5self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))elif os.path.isfile(source):self.type = 'img'elif os.path.isdir(source):img_files = ['%s/%s' % (i[0].replace("\\", "/"), j) for i in os.walk(source) for j in i[-1] if j.endswith(('jpg', 'png', 'jpeg', 'JPG'))]self.img_files= natsorted(img_files)self.img_index = 0self.type = 'dir_img'self.total_frames = len(self.img_files)self.source = sourcedef get_filename(self):if self.suffix != 'mp4':if self.type == 'dir':if self.img_index < 0 or self.img_index >= len(self.img_files):return None,None,Noneimg_file = self.img_files[self.img_index]self.img_index += 1return img_file, self.img_index, img_fileelse:if self.img_index==1:return None,None,Noneself.img_index += 1return self.source, 0, Nonedef get_img(self):if self.type == 'mp4':ret, frame = self.cap.read()if not ret:self.cap.release()return None,None,Noneself.img_index += 1return frame,self.img_index,Noneelif self.type == 'img':return cv2.imread(self.source),0,Noneelif self.type == 'dir_img':if self.img_index < 0 or self.img_index >= len(self.img_files):return None,None,Noneimg_file = self.img_files[self.img_index]img = cv2.imread(img_file)self.img_index+=1return img,self.img_index,img_fileelse:return None,None,None