ImageFolderVideoReader.py

import cv2
import os
import glob
import numpy as np
from pathlib import Path

class ImageFolderVideoReader(object):
    """
    A class to read a folder of images as if it were a video file using OpenCV.
    """
    
    def __init__(self, folder_path, image_extensions=None):
        """
        Initialize the ImageFolderVideoReader.
        
        Args:
            folder_path (str): Path to the folder containing images
            image_extensions (list): List of image extensions to include
        """
        if image_extensions is None:
            image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
        
        self.folder_path = Path(folder_path)
        self.image_extensions = [ext.lower() for ext in image_extensions]
        self.image_paths = self._get_image_paths()
        self.current_frame = 0
        self.total_frames = len(self.image_paths)
        
        if self.total_frames == 0:
            raise ValueError(f"No images found in {folder_path}")
        
        # Read first image to get dimensions
        first_image = cv2.imread(str(self.image_paths[0]))
        if first_image is None:
            raise ValueError(f"Cannot read first image: {self.image_paths[0]}")
        
        self.frame_height, self.frame_width = first_image.shape[:2]
        self.fps = 30  # Default FPS, can be changed
    
    def _get_image_paths(self):
        """Get sorted list of image paths in the folder."""
        image_paths = []
        for ext in self.image_extensions:
            pattern = str(self.folder_path / f"*{ext}")
            image_paths.extend(glob.glob(pattern))
            pattern = str(self.folder_path / f"*{ext.upper()}")
            image_paths.extend(glob.glob(pattern))
        
        # Remove duplicates and sort
        image_paths = list(set(image_paths))
        image_paths.sort()  # Natural sorting might be needed for numbered files
        
        return [Path(p) for p in image_paths]
    
    def read(self):
        """
        Read the next frame from the image sequence.
        
        Returns:
            tuple: (success, frame) where success is bool and frame is numpy array
        """
        if self.current_frame >= self.total_frames:
            return False, None
        
        image_path = self.image_paths[self.current_frame]
        frame = cv2.imread(str(image_path))
        
        if frame is None:
            print(f"Warning: Could not read image {image_path}")
            self.current_frame += 1
            return self.read()  # Try next image
        
        self.current_frame += 1
        return True, frame
    
    def get(self, prop):
        """
        Get video properties (similar to cv2.VideoCapture.get).
        
        Args:
            prop: Property ID (use cv2.CAP_PROP_* constants)
        
        Returns:
            Property value
        """
        if prop == cv2.CAP_PROP_FRAME_COUNT:
            return self.total_frames
        elif prop == cv2.CAP_PROP_FRAME_WIDTH:
            return self.frame_width
        elif prop == cv2.CAP_PROP_FRAME_HEIGHT:
            return self.frame_height
        elif prop == cv2.CAP_PROP_FPS:
            return self.fps
        elif prop == cv2.CAP_PROP_POS_FRAMES:
            return self.current_frame
        else:
            return 0.0
    
    def set(self, prop, value):
        """
        Set video properties (similar to cv2.VideoCapture.set).
        
        Args:
            prop: Property ID
            value: New value
        
        Returns:
            bool: Success status
        """
        if prop == cv2.CAP_PROP_POS_FRAMES:
            if 0 <= value < self.total_frames:
                self.current_frame = int(value)
                return True
        elif prop == cv2.CAP_PROP_FPS:
            self.fps = value
            return True
        return False
    
    def release(self):
        """Release resources (for compatibility with cv2.VideoCapture)."""
        pass
    
    def isOpened(self):
        """Check if the image folder is successfully opened."""
        return self.total_frames > 0