train_image_model.py

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)


DATA_DIR = "data/images"   # ✅ YE HI SAHI HAI

print("Dataset path:", os.path.abspath(DATA_DIR))

if not os.path.exists(DATA_DIR):
    print("❌ data/images folder not found!")
    exit()


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


dataset = datasets.ImageFolder(root=DATA_DIR, transform=transform)

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

loader = DataLoader(dataset, batch_size=4, shuffle=True)


model = models.resnet18(weights="IMAGENET1K_V1")
model.fc = nn.Linear(model.fc.in_features, len(dataset.classes))
model = model.to(device)

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


epochs = 3

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}")


os.makedirs("models", exist_ok=True)
torch.save(model.state_dict(), "models/image_model.pth")

print("✅ Model training complete & saved!")