72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
|
import os
|
||
|
|
import cv2
|
||
|
|
import torch
|
||
|
|
import numpy as np
|
||
|
|
from torch.utils.data import DataLoader, TensorDataset
|
||
|
|
from torch import nn, optim
|
||
|
|
from BN import CatDogClassifier
|
||
|
|
import time
|
||
|
|
|
||
|
|
# Define image size
|
||
|
|
IMG_SIZE = 128
|
||
|
|
|
||
|
|
# Function to load images from a directory
|
||
|
|
def load_images_from_folder(folder, label):
|
||
|
|
data = []
|
||
|
|
for filename in os.listdir(folder):
|
||
|
|
img_path = os.path.join(folder, filename)
|
||
|
|
img = cv2.imread(img_path) # Load image
|
||
|
|
if img is not None:
|
||
|
|
img = cv2.resize(img, (IMG_SIZE, IMG_SIZE)) # Resize
|
||
|
|
img = img / 255.0 # Normalize
|
||
|
|
data.append((img, label))
|
||
|
|
return data
|
||
|
|
|
||
|
|
# Load dataset
|
||
|
|
cat_data = load_images_from_folder("dataset/training_set/cats", label=0) # Label 0 for cats
|
||
|
|
dog_data = load_images_from_folder("dataset/training_set/dogs", label=1) # Label 1 for dogs
|
||
|
|
|
||
|
|
# Combine and shuffle data
|
||
|
|
dataset = cat_data + dog_data
|
||
|
|
np.random.shuffle(dataset)
|
||
|
|
|
||
|
|
# Convert to NumPy arrays
|
||
|
|
X = np.array([item[0] for item in dataset], dtype=np.float32) # Image data
|
||
|
|
Y = np.array([item[1] for item in dataset], dtype=np.int64) # Labels
|
||
|
|
|
||
|
|
# Reshape to PyTorch format (N, C, H, W)
|
||
|
|
X = np.transpose(X, (0, 3, 1, 2)) # Convert to (batch, channels, height, width)
|
||
|
|
|
||
|
|
# Convert to PyTorch tensors
|
||
|
|
X_tensor = torch.tensor(X).to("cuda")
|
||
|
|
Y_tensor = torch.tensor(Y).to("cuda")
|
||
|
|
|
||
|
|
# Create dataset and data loader
|
||
|
|
dataset = TensorDataset(X_tensor, Y_tensor)
|
||
|
|
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
|
||
|
|
|
||
|
|
# Initialize model, loss, and optimizer
|
||
|
|
model = CatDogClassifier()
|
||
|
|
model = model.to("cuda")
|
||
|
|
criterion = nn.CrossEntropyLoss()
|
||
|
|
optimizer = optim.Adam(model.parameters(), lr=0.001)
|
||
|
|
|
||
|
|
# Training loop
|
||
|
|
num_epochs = 25
|
||
|
|
start_time = time.time()
|
||
|
|
for epoch in range(num_epochs):
|
||
|
|
total_loss = 0
|
||
|
|
for images, labels in dataloader:
|
||
|
|
optimizer.zero_grad() # Reset gradients
|
||
|
|
outputs = model(images) # Forward pass
|
||
|
|
loss = criterion(outputs, labels) # Compute loss
|
||
|
|
loss.backward() # Backpropagation
|
||
|
|
optimizer.step() # Update weights
|
||
|
|
total_loss += loss.item()
|
||
|
|
|
||
|
|
print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(dataloader):.4f}")
|
||
|
|
|
||
|
|
print(f"Time taken: {(time.time()-start_time):.2f} seconds")
|
||
|
|
|
||
|
|
|
||
|
|
torch.save(model, "super_model.pth")
|