commit a271c738d791a44b19d814f8c278940841c4e42d Author: Marvin Kraußer Date: Tue Apr 28 22:41:12 2026 +0200 first commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..cdd2532 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +bird_cnn/data/ \ No newline at end of file diff --git a/bird_cnn/__pycache__/bird_cnn.cpython-314.pyc b/bird_cnn/__pycache__/bird_cnn.cpython-314.pyc new file mode 100644 index 0000000..51efc66 Binary files /dev/null and b/bird_cnn/__pycache__/bird_cnn.cpython-314.pyc differ diff --git a/bird_cnn/bird_cnn.py b/bird_cnn/bird_cnn.py new file mode 100644 index 0000000..751789c --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/bird_cnn/main.py b/bird_cnn/main.py new file mode 100644 index 0000000..6d3802d --- /dev/null +++ b/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}") \ No newline at end of file diff --git a/bird_cnn/saved_models/bird_cnn/bird_cnn b/bird_cnn/saved_models/bird_cnn/bird_cnn new file mode 100644 index 0000000..a24da97 Binary files /dev/null and b/bird_cnn/saved_models/bird_cnn/bird_cnn differ diff --git a/bird_cnn/test.jpg b/bird_cnn/test.jpg new file mode 100644 index 0000000..2df10da Binary files /dev/null and b/bird_cnn/test.jpg differ