Spaces:
Running
Running
| from PIL import Image | |
| import torch | |
| import numpy as np | |
| from typing import IO | |
| import cv2 | |
| from torchvision import transforms | |
| # Import the globally loaded models instance | |
| from model_loader import models | |
| class ImagePreprocessor: | |
| """ | |
| Handles preprocessing of images for the FFT CNN model. | |
| """ | |
| def __init__(self): | |
| """ | |
| Initializes the preprocessor. | |
| """ | |
| self.device = models.device | |
| # Define the image transformations, matching the training process | |
| self.transform = transforms.Compose([ | |
| transforms.ToPILImage(), | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| ]) | |
| def process(self, image_file: IO) -> torch.Tensor: | |
| """ | |
| Opens an image file, applies FFT, preprocesses it, and returns a tensor. | |
| Args: | |
| image_file (IO): The image file object (e.g., from a file upload). | |
| Returns: | |
| torch.Tensor: The preprocessed image as a tensor, ready for the model. | |
| """ | |
| try: | |
| # Read the image file into a numpy array | |
| image_np = np.frombuffer(image_file.read(), np.uint8) | |
| # Decode the image as grayscale | |
| img = cv2.imdecode(image_np, cv2.IMREAD_GRAYSCALE) | |
| except Exception as e: | |
| print(f"Error reading or decoding image: {e}") | |
| raise ValueError("Invalid or corrupted image file.") | |
| if img is None: | |
| raise ValueError("Could not decode image. File may be empty or corrupted.") | |
| # 1. Apply Fast Fourier Transform (FFT) | |
| f = np.fft.fft2(img) | |
| fshift = np.fft.fftshift(f) | |
| magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1) # Add 1 to avoid log(0) | |
| # Normalize the magnitude spectrum to be in the range [0, 255] | |
| magnitude_spectrum = cv2.normalize(magnitude_spectrum, None, 0, 255, cv2.NORM_MINMAX) | |
| magnitude_spectrum = np.uint8(magnitude_spectrum) | |
| # 2. Apply torchvision transforms | |
| image_tensor = self.transform(magnitude_spectrum) | |
| # Add a batch dimension and move to the correct device | |
| image_tensor = image_tensor.unsqueeze(0).to(self.device) | |
| return image_tensor | |
| # Create a single instance of the preprocessor | |
| preprocessor = ImagePreprocessor() | |