Upload 4 files

script/predict.py

@@ -0,0 +1,33 @@
+import sys
+import os
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+import torch
+from PIL import Image
+from torchvision import transforms
+from models.cnn_model import CatBreedCNN
+
+
+classes = ['Bengal', 'Domestic_Shorthair', 'Maine_Coon','Ragdoll','Siamese',]  # Update as needed
+
+def predict(image_path):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    model = CatBreedCNN(num_classes=len(classes))
+    model.load_state_dict(torch.load("models/cat_cnn.pth", map_location=device))
+    model.eval()
+    
+    transform = transforms.Compose([
+        transforms.Resize((128, 128)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5]*3, [0.5]*3)
+    ])
+    
+    image = Image.open(image_path).convert("RGB")
+    image = transform(image).unsqueeze(0).to(device)
+    
+    with torch.no_grad():
+        output = model(image)
+        predicted_index = output.argmax(dim=1).item()
+    
+    return classes[predicted_index]

script/train.py

@@ -0,0 +1,37 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from models.cnn_model import CatBreedCNN
+from utils.data_loader import get_dataloaders
+from utils.evaluate import evaluate_model
+
+# Load data
+train_loader, val_loader, classes = get_dataloaders("data/cat_breed_dataset")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Initialize model
+model = CatBreedCNN(len(classes)).to(device)
+
+# Loss and optimizer
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+# Training loop
+for epoch in range(20):
+    model.train()
+    for x, y in train_loader:
+        x, y = x.to(device), y.to(device)
+        optimizer.zero_grad()
+        outputs = model(x)
+        loss = criterion(outputs, y)
+        loss.backward()
+        optimizer.step()
+    
+    print(f"Epoch {epoch+1} complete. Evaluating...")  # ✅ Corrected f-string
+
+    # Evaluate
+    report, _ = evaluate_model(model, val_loader, device)
+    print(report)
+
+# Save model
+torch.save(model.state_dict(), "models/cat_cnn.pth")