diff --git a/bird_cnn/__pycache__/bird_cnn.cpython-314.pyc b/bird_cnn/__pycache__/bird_cnn.cpython-314.pyc index 51efc66..040f7ad 100644 Binary files a/bird_cnn/__pycache__/bird_cnn.cpython-314.pyc and b/bird_cnn/__pycache__/bird_cnn.cpython-314.pyc differ diff --git a/bird_cnn/__pycache__/main.cpython-314.pyc b/bird_cnn/__pycache__/main.cpython-314.pyc new file mode 100644 index 0000000..f953b57 Binary files /dev/null and b/bird_cnn/__pycache__/main.cpython-314.pyc differ diff --git a/bird_cnn/__pycache__/server.cpython-314.pyc b/bird_cnn/__pycache__/server.cpython-314.pyc new file mode 100644 index 0000000..65954ff Binary files /dev/null and b/bird_cnn/__pycache__/server.cpython-314.pyc differ diff --git a/bird_cnn/bird_cnn.py b/bird_cnn/bird_cnn.py index 751789c..a43a205 100644 --- a/bird_cnn/bird_cnn.py +++ b/bird_cnn/bird_cnn.py @@ -16,13 +16,14 @@ class Bird_CNN(nn.Module): 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): diff --git a/bird_cnn/main.py b/bird_cnn/main.py index 6d3802d..7ea43d9 100644 --- a/bird_cnn/main.py +++ b/bird_cnn/main.py @@ -10,6 +10,7 @@ from enum import Enum from PIL import Image SAVE_PATH = "./saved_models" +#IMAGE_SIZE = (1141, 850) IMAGE_SIZE = (300, 300) class bird_species(Enum): @@ -28,11 +29,10 @@ transform = transforms.Compose([ 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_size = int(0.8 * len(dataset)) +val_size = len(dataset) - train_size -train_dataset, val_dataset, _ = random_split(dataset, [train_size, val_size, throw_away]) +train_dataset, val_dataset = random_split(dataset, [train_size, val_size]) train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False) @@ -48,7 +48,7 @@ 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_path = "test1.jpg" image = Image.open(image_path).convert("RGB") image = transform(image) diff --git a/bird_cnn/saved_models/bird_cnn/bird_cnn b/bird_cnn/saved_models/bird_cnn/bird_cnn index a24da97..b672e31 100644 Binary files a/bird_cnn/saved_models/bird_cnn/bird_cnn and b/bird_cnn/saved_models/bird_cnn/bird_cnn differ diff --git a/bird_cnn/server.py b/bird_cnn/server.py new file mode 100644 index 0000000..646ab35 --- /dev/null +++ b/bird_cnn/server.py @@ -0,0 +1,61 @@ +from enum import Enum +import os + +import torch +from torchvision import transforms +from fastapi import FastAPI, File, UploadFile +from PIL import Image +import io +import torch.nn.functional as F + +from bird_cnn import Bird_CNN + +SAVE_PATH = "./saved_models" +#IMAGE_SIZE = (1141, 850) +IMAGE_SIZE = (300, 300) + +transform = transforms.Compose([ + transforms.Resize(IMAGE_SIZE), + transforms.ToTensor() +]) + +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() +]) + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +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]) +full_path = os.path.join(SAVE_PATH, "bird_cnn", "bird_cnn") +model.load_state_dict(torch.load(full_path, weights_only=False)) +model.to(device) +model.eval() + +app = FastAPI() + +@app.post("/predict") +async def predict(file: UploadFile = File(...)): + image_bytes = await file.read() + + image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + image = transform(image).unsqueeze(0).to(device) + + with torch.no_grad(): + pred = model(image) + probs = F.softmax(pred, dim=1) + confidence, cls = torch.max(probs, dim=1) + + return { + "class": bird_species(cls.item()).name, + "confidence": confidence.item() + } \ No newline at end of file diff --git a/bird_cnn/test1.jpg b/bird_cnn/test1.jpg new file mode 100644 index 0000000..512da5a Binary files /dev/null and b/bird_cnn/test1.jpg differ diff --git a/bird_cnn/util.py b/bird_cnn/util.py new file mode 100644 index 0000000..5aab5d2 --- /dev/null +++ b/bird_cnn/util.py @@ -0,0 +1,22 @@ +from PIL import Image +import os + +folder = "data/train" + +widths = [] +heights = [] + +for root, _, files in os.walk(folder): + for file in files: + if file.endswith((".jpg", ".png", ".jpeg")): + path = os.path.join(root, file) + img = Image.open(path) + w, h = img.size + widths.append(w) + heights.append(h) + +avg_w = sum(widths) / len(widths) +avg_h = sum(heights) / len(heights) + +print("Average width:", avg_w) +print("Average height:", avg_h) \ No newline at end of file