Spaces:
Runtime error
Runtime error
| import os | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torch.fft | |
| import torchvision.models as models | |
| from torch.utils.data import Dataset | |
| from PIL import Image | |
| def extract_frequency_features(image_tensor): | |
| # Convert RGB to Grayscale | |
| gray_image = 0.2989 * image_tensor[0] + 0.5870 * image_tensor[1] + 0.1140 * image_tensor[2] | |
| # 2D Fast Fourier Transform to look at noise grids | |
| fft_result = torch.fft.fft2(gray_image) | |
| fft_shifted = torch.fft.fftshift(fft_result) | |
| magnitude_spectrum = torch.log(torch.abs(fft_shifted) + 1e-8) | |
| return magnitude_spectrum.unsqueeze(0) | |
| class BluemarkDataset(Dataset): | |
| def __init__(self, data_dir, transform=None): | |
| self.data_dir = data_dir | |
| self.transform = transform | |
| self.image_paths = [] | |
| self.labels = [] | |
| self.classes = {'clean': 0, 'marked': 1, 'scrubbed': 2} | |
| for class_name, label in self.classes.items(): | |
| class_dir = os.path.join(data_dir, class_name) | |
| if os.path.isdir(class_dir): | |
| for img_name in os.listdir(class_dir): | |
| if img_name.lower().endswith(('.png', '.jpg', '.jpeg')): | |
| self.image_paths.append(os.path.join(class_dir, img_name)) | |
| self.labels.append(label) | |
| def __len__(self): | |
| return len(self.image_paths) | |
| def __getitem__(self, idx): | |
| img_path = self.image_paths[idx] | |
| image = Image.open(img_path).convert('RGB') | |
| label = self.labels[idx] | |
| if self.transform: | |
| image = self.transform(image) | |
| freq_features = extract_frequency_features(image) | |
| return freq_features, label | |
| class ResNet50FrequencyDetector(nn.Module): | |
| def __init__(self, num_classes=3): | |
| super().__init__() | |
| # Load the massive ResNet-50 architecture | |
| self.resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) | |
| # SURGICAL MODIFICATION: | |
| # ResNet expects 3 RGB channels. We change the first layer to accept our 1-channel FFT map. | |
| original_conv1 = self.resnet.conv1 | |
| self.resnet.conv1 = nn.Conv2d( | |
| 1, # Changed from 3 to 1 | |
| original_conv1.out_channels, | |
| kernel_size=original_conv1.kernel_size, | |
| stride=original_conv1.stride, | |
| padding=original_conv1.padding, | |
| bias=False | |
| ) | |
| # Modify the final classification layer to output our 3 classes instead of 1000 | |
| num_ftrs = self.resnet.fc.in_features | |
| self.resnet.fc = nn.Sequential( | |
| nn.Dropout(0.5), # Keep dropout to prevent memorizing the dataset | |
| nn.Linear(num_ftrs, num_classes) | |
| ) | |
| def forward(self, x): | |
| return self.resnet(x) |