import torch
import torch.nn as nn
from torchvision import transforms
from torchvision.models import resnet50
from PIL import Image

class ImageNormalizer(nn.Module):
    def __init__(self, mean, std):
        super(ImageNormalizer, self).__init__()
        self.mean = torch.Tensor(mean)
        self.std = torch.Tensor(std)

    def forward(self, x):
        return (x - self.mean.type_as(x)[None, :, None, None]) / self.std.type_as(x)[None, :, None, None]

class DeltaEvaluator:
    def __init__(self, imagenet_class_index, orig_image):
        self.image_normalizer = ImageNormalizer(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        self.preprocess = transforms.Compose([
            transforms.Resize(224),
            transforms.ToTensor(),
        ])
        with open(imagenet_class_index) as f:
            self.imagenet_classes = {int(i): x[1] for i, x in json.load(f).items()}
        orig_img = Image.open(orig_image)
        self.image_tensor = self.preprocess(orig_img)[None, :, :, :]

    def check_delta(self, resnet50_path, delta_path):
        model = resnet50(weights=None)
        model.load_state_dict(torch.load(resnet50_path))
        model.eval()

        delta = torch.load(delta_path)
        prediction = model(self.image_normalizer(self.image_tensor + delta))

        max_class = prediction.max(dim=1)[1].item()
        predicted_class = self.imagenet_classes[max_class], max_class
        probability = torch.nn.Softmax(dim=1)(prediction)[0, max_class].item()

        return predicted_class[0], probability

def generate_delta(image_path, target_class_index):
    model = resnet50(pretrained=True)
    model.eval()

    preprocess = transforms.Compose([
        transforms.Resize(224),
        transforms.ToTensor(),
    ])

    orig_image = Image.open(image_path)
    image_tensor = preprocess(orig_image)[None, :, :, :]

    mean = torch.Tensor([0.485, 0.456, 0.406])
    std = torch.Tensor([0.229, 0.224, 0.225])
    image_normalizer = lambda x: (x - mean.type_as(x)[None, :, None, None]) / std.type_as(x)[None, :, None, None]
    normalized_image = image_normalizer(image_tensor)

    target_class = torch.tensor([target_class_index])

    criterion = nn.CrossEntropyLoss()
    delta = torch.zeros_like(image_tensor, requires_grad=True)
    optimizer = torch.optim.Adam([delta], lr=0.01)

    for step in range(100):
        optimizer.zero_grad()
        output = model(image_normalizer(image_tensor + delta))
        loss = criterion(output, target_class)
        loss.backward()
        optimizer.step()

    torch.save(delta, 'delta.pt')

generate_delta('resim.jpg', 510)

evaluator = DeltaEvaluator('imagenet_class_index.json', 'resim.jpg')
predicted_class, probability = evaluator.check_delta('resnet50_weights.pth', 'delta.pt')

print(f"Tahmin edilen sınıf: {predicted_class}")