ai_generated_content_detector/train_image_model_auto.py at main · ishasinghrathore37/ai_generated_content_detector · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
import os
import shutil
from PIL import Image
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.models import resnet18


SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = SCRIPT_DIR  # Script folder ko root maan lo

DATA_DIR = os.path.join(PROJECT_ROOT, "data", "images")
REAL_DIR = os.path.join(DATA_DIR, "real")
AI_DIR = os.path.join(DATA_DIR, "ai")

MODEL_DIR = os.path.join(PROJECT_ROOT, "models", "image_model")
os.makedirs(MODEL_DIR, exist_ok=True)


real_image_path = r"C:\Users\Isha\Pictures\Saved Pictures\image1.jpg"
ai_image_path   = r"C:\Users\Isha\Pictures\Saved Pictures\image2.jpg"


os.makedirs(REAL_DIR, exist_ok=True)
os.makedirs(AI_DIR, exist_ok=True)

# Copy image safely
shutil.copy(real_image_path, REAL_DIR)
shutil.copy(ai_image_path, AI_DIR)

print("✅ Images copied to folders successfully")


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

dataset = datasets.ImageFolder(root=DATA_DIR, transform=transform)
loader = DataLoader(dataset, batch_size=2, shuffle=True)

print("Classes found:", dataset.classes)

model = resnet18(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, 2)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

epochs = 2
for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f"Epoch [{epoch+1}/{epochs}] Loss: {running_loss:.4f}")


model_path = os.path.join(MODEL_DIR, "image_ai_detector.pth")
torch.save(model.state_dict(), model_path)
print("✅ Model saved at:", model_path)