Update preprocess.py

preprocess.py

@@ -1,33 +1,39 @@
-import cv2
-import numpy as np
-
-# Load OpenCV's built-in face detector
-face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-
-def extract_face(image_path, output_size=(224, 224)):
-    image = cv2.imread(image_path)
-    if image is None:
-        return None
-        
-    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    
-    # Detect faces
-    faces = face_cascade.detectMultiScale(gray_image, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
-
-    if len(faces) > 0:
-        # Grab the coordinates of the first face found
-        (x, y, w, h) = faces[0]
-        ih, iw, _ = image.shape
-        
-        # Add a small margin around the face
-        face_crop = image[max(0, y-20):min(ih, y+h+20), max(0, x-20):min(iw, x+w+20)]
-        if face_crop.size == 0:
-            return None
-            
-        # Resize for the neural network
-        face_resized = cv2.resize(face_crop, output_size)
-        # Convert to grayscale for FFT
-        face_gray = cv2.cvtColor(face_resized, cv2.COLOR_BGR2GRAY)
-        return face_gray
-        
-    return None
+import cv2
+import numpy as np
+import torch
+from torchvision import transforms
+from PIL import Image
+import io
+
+def prepare_tensor(file_bytes):
+    # 1. Load image in grayscale
+    nparr = np.frombuffer(file_bytes, np.uint8)
+    img = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE)
+    
+    # 2. THE FIX: Micro-Blur instead of Heavy Blur
+    # Reduced from (5,5) to (3,3) to preserve GAN/Diffusion high-frequency artifacts
+    img_smoothed = cv2.GaussianBlur(img, (3, 3), 0)
+    
+    # 3. 2D Fast Fourier Transform (FFT)
+    f = np.fft.fft2(img_smoothed)
+    fshift = np.fft.fftshift(f)
+    
+    # 4. Enhance the Magnitude Spectrum
+    magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1e-8)
+    
+    # 5. Normalize for ResNet-18 (0-255)
+    magnitude_spectrum = cv2.normalize(magnitude_spectrum, None, 0, 255, cv2.NORM_MINMAX)
+    magnitude_spectrum = np.uint8(magnitude_spectrum)
+    
+    # Convert back to 3-channel RGB as ResNet expects 3 channels
+    img_rgb = cv2.cvtColor(magnitude_spectrum, cv2.COLOR_GRAY2RGB)
+    
+    # 6. Final PyTorch Tensor Transformations
+    transform = transforms.Compose([
+        transforms.ToPILImage(),
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    
+    return transform(img_rgb).unsqueeze(0)