first commit

.gitignore

@@ -0,0 +1 @@
+bird_cnn/data/

bird_cnn/bird_cnn.py

@@ -0,0 +1,97 @@
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from tqdm import tqdm
+
+
+class Bird_CNN(nn.Module):
+    def __init__(self, c_in, c_hidden, c_out, kernel_size, img_width, img_height):
+        super().__init__()
+        self.model = nn.Sequential(
+            nn.Conv2d(c_in, c_hidden, kernel_size, padding=kernel_size//2),
+            nn.ReLU(),
+
+            nn.Conv2d(c_hidden, c_hidden, kernel_size, padding=kernel_size//2),
+            nn.ReLU(),
+            nn.Flatten(),
+            nn.Linear(c_hidden * img_height * img_width, c_out)
+        )
+
+    def forward(self, x):
+        return self.model(x)
+    
+
+def trainCNN(model, optimizer, loss_module, train_data_loader, validation_data_loader, device, num_epochs, SAVE_PATH, save=False):
+
+    best_val = 0
+    for epoch in range(num_epochs):
+        ############
+        # Training #
+        ############
+        model.train()
+
+        true_preds, count = 0, 0
+        for data_inputs, classes in tqdm(train_data_loader, desc=f"Train Epoch {epoch+1}", leave=False):
+            data_inputs = data_inputs.to(device)
+            classes = classes.to(device)
+
+            preds = model(data_inputs)
+
+            loss = loss_module(preds, classes)
+
+            optimizer.zero_grad()
+
+            loss.backward()
+
+            optimizer.step()
+
+            true_preds += (preds.argmax(dim=1) == classes).sum().item()
+            count += data_inputs.size(0)
+        train_acc = true_preds / count
+
+        torch.cuda.empty_cache()
+
+        ##############
+        # Validation #
+        ##############
+        model.eval()
+
+        true_preds, count = 0, 0
+        for data_inputs, classes in tqdm(validation_data_loader, desc=f"Validate Epoch {epoch+1}", leave=False):
+            with torch.no_grad():
+                data_inputs = data_inputs.to(device)
+                classes = classes.to(device)
+
+                preds = model(data_inputs)
+
+                loss = loss_module(preds, classes)
+
+                true_preds += (preds.argmax(dim=1) == classes).sum().item()
+                count += data_inputs.size(0)
+        val_acc = true_preds / count
+
+        if(save and best_val < val_acc):
+            best_val = val_acc
+            save_dir = os.path.join(SAVE_PATH, "bird_cnn")
+            os.makedirs(save_dir, exist_ok=True)
+
+            save_path = os.path.join(save_dir, "bird_cnn")
+            torch.save(model.state_dict(), save_path)
+
+        print(f"epoch: {epoch+1} | train accuracy: {int(train_acc * 1000) / 10}% | validation accuracy: {int(val_acc * 1000) / 10}%")
+        torch.cuda.empty_cache()
+
+
+def sample(model, img, device, SAVE_PATH, model_name="bird_cnn", folder="bird_cnn"):
+    with torch.no_grad():
+        full_path = os.path.join(SAVE_PATH, folder, model_name)
+        state_dict = torch.load(full_path, weights_only=False)
+        model.load_state_dict(state_dict)
+        model.eval()
+        img = img.to(device)
+        pred = model(img)
+        probs = F.softmax(pred, dim=1)
+        return(torch.max(probs, dim=1))

bird_cnn/main.py

@@ -0,0 +1,58 @@
+import torch
+import torch.nn as nn
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader, random_split
+
+from bird_cnn import Bird_CNN, sample, trainCNN
+
+from enum import Enum
+
+from PIL import Image
+
+SAVE_PATH = "./saved_models"
+IMAGE_SIZE = (300, 300)
+
+class bird_species(Enum):
+    Common_Kingfisher = 0
+    CommonMyna = 1
+    House_Crow = 2
+    Indian_Peacock = 3
+    Indian_Pitta = 4
+    Ruddy_Shelduck = 5
+    Sarus_Crane = 6
+
+transform = transforms.Compose([
+    transforms.Resize(IMAGE_SIZE),
+    transforms.ToTensor()
+])
+
+dataset = datasets.ImageFolder("data/train", transform=transform)
+
+train_size = int(0.3 * len(dataset))
+val_size = int(len(dataset) * 0.3)
+throw_away = len(dataset) - val_size - train_size
+
+train_dataset, val_dataset, _ = random_split(dataset, [train_size, val_size, throw_away])
+
+train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
+
+device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
+print("Using device", device)
+
+model = Bird_CNN(c_in=3, c_hidden=15, c_out=7, kernel_size=3, img_width=IMAGE_SIZE[0], img_height=IMAGE_SIZE[1])
+model.to(device)
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
+loss_module = nn.CrossEntropyLoss()
+
+trainCNN(model, optimizer, loss_module, train_loader, val_loader, device, 50, SAVE_PATH=SAVE_PATH, save=True)
+exit()
+
+image_path = "test.jpg"
+
+image = Image.open(image_path).convert("RGB")
+image = transform(image)
+image = image.unsqueeze(0)
+
+confidence, pred = sample(model=model, img=image, device=device, SAVE_PATH=SAVE_PATH)
+print(f"Species: {bird_species(pred.item()).name} | Confidence: {int(confidence.item()*100)/100}")