image/image_base.py

import os
import json
from datetime import datetime, timedelta
from process_exception import ProcessAlreadyRunning
from main_base import MainBase
from pathlib import Path

class ImageBase(MainBase):
	def __init__(self, feature=None):
		super().__init__()
		self.current_path = os.path.abspath(os.path.join(os.path.dirname(__file__)))
		self.feature = feature
		self.input_dir = os.path.join(self.current_path, "input")
		self.output_dir = os.path.join(self.current_path, "output")
		self.supported_formats = {
			'jpg': 'JPEG',
			'jpeg': 'JPEG', 
			'png': 'PNG',
			'webp': 'WEBP',
			'tiff': 'TIFF',
			'tif': 'TIFF',
			'bmp': 'BMP',
			'gif': 'GIF',
			'heic': 'HEIC'
		}

		# Create directories if they don't exist
		os.makedirs(self.input_dir, exist_ok=True)
		os.makedirs(self.output_dir, exist_ok=True)

	def _validate_input_file(self, image=None):
		"""
		Validate that the input file exists and is a supported image format.
		Args:
			image: UploadFile object (for upload case)
		Raises:
			FileNotFoundError: If input file doesn't exist
			ValueError: If file format is not supported
		"""
		if image:
			# Handle upload case
			if not image.filename:
				raise ValueError("Uploaded file must have a filename")
			
			extension = Path(image.filename).suffix.lower().lstrip('.')
		else:
			# Handle file path case
			if not self.input_file_path or not self.input_file_path.strip():
				raise ValueError("Image path cannot be empty")
			
			path_obj = Path(self.input_file_path)
			if not path_obj.exists():
				raise FileNotFoundError(f"Image file not found: {self.input_file_path}")
			
			if not path_obj.is_file():
				raise ValueError(f"Path is not a file: {self.input_file_path}")
			
			extension = path_obj.suffix.lower().lstrip('.')

		# Check if extension is supported
		if extension not in self.supported_formats:
			supported_list = ', '.join(self.supported_formats.keys())
			raise ValueError(f"Unsupported format '{extension}'. Supported: {supported_list}")

	def upload_validate(self, image):
		super().upload_validate(image)
		self._validate_input_file(image=image)