Update README.md

README.md

@@ -44,36 +44,35 @@ from transformers import AutoModelForImageClassification
 
 def preprocess_for_artifact_detection(image_path):
     """
-    Loads an STM image and converts it to the required 3-channel format 
-    (grayscale, FFT amplitude, FFT phase) for the model.
+    Loads an STM image and converts it to the required 3-channel format
+    (grayscale, magnitude spectrum, phase) for the model.
     """
-    # 1. Load and prepare grayscale channel
-    with Image.open(image_path) as img:
-        img = img.convert('L').resize((224, 224))
-        grayscale_img = np.array(img) / 255.0
-
-    # 2. Compute FFT, Amplitude, and Phase
+    try:
+        with Image.open(image_path) as img:
+            img = img.convert('L').resize((224, 224))
+            grayscale_img = np.array(img) / 255.0
+    except FileNotFoundError:
+        print(f"Error: The file at {image_path} was not found.")
+        return None
+
+    # Compute FFT, Magnitude, and Phase 
     fft_data = np.fft.fft2(grayscale_img)
     fft_shifted = np.fft.fftshift(fft_data)
     
-    amplitude = np.log1p(np.abs(fft_shifted))
+    magnitude_spectrum = np.abs(fft_shifted)
     phase = np.angle(fft_shifted)
 
-    # 3. Normalize channels to be in a 0-1 range
-    amplitude = (amplitude - np.min(amplitude)) / (np.max(amplitude) - np.min(amplitude))
-    phase = (phase - np.min(phase)) / (np.max(phase) - np.min(phase))
-    
-    # 4. Stack channels and convert to PyTorch tensor (C, H, W)
-    stacked_channels = np.stack([grayscale_img, amplitude, phase], axis=0)
+    # Stack channels and convert to PyTorch tensor (C, H, W)
+    stacked_channels = np.stack([grayscale_img, magnitude_spectrum, phase], axis=0)
     
-    # 5. Add a batch dimension (B, C, H, W) and return as float tensor
+    # Add a batch dimension (B, C, H, W) and return as float tensor
     return torch.tensor(stacked_channels, dtype=torch.float32).unsqueeze(0)
 
 # Load the model from the Hub
 model_name = "t0m-R/vit-stm-artifact-fft"
 model = AutoModelForImageClassification.from_pretrained(model_name)
 
-# Preprocess your image
+# Preprocess 
 image_path = "path/to/your/stm_image" # Replace with your image path
 preprocessed_image = preprocess_for_artifact_detection(image_path)