changed website structure, added new cnn files
This commit is contained in:
@@ -0,0 +1,51 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision
|
||||
from torchvision import datasets, transforms
|
||||
from torch.utils.data import DataLoader, random_split
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
from bird_cnn import Bird_CNN, sample, trainCNN
|
||||
|
||||
from enum import Enum
|
||||
|
||||
from PIL import Image
|
||||
|
||||
SAVE_PATH = "./saved_models"
|
||||
IMAGE_SIZE = 64
|
||||
|
||||
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.CenterCrop(IMAGE_SIZE),
|
||||
transforms.ToTensor()
|
||||
])
|
||||
|
||||
dataset = datasets.ImageFolder("data/train", transform=transform)
|
||||
|
||||
train_size = int(0.8 * len(dataset))
|
||||
val_size = len(dataset) - train_size
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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=32, c_out=7)
|
||||
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)
|
||||
@@ -0,0 +1,85 @@
|
||||
from collections import defaultdict
|
||||
from enum import Enum
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
import cv2
|
||||
|
||||
import torch
|
||||
from torchvision import transforms
|
||||
import torch.nn.functional as F
|
||||
|
||||
from bird_cnn import Bird_CNN
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
|
||||
BUILD_PATH = "./build_models"
|
||||
IMAGE_SIZE = 64
|
||||
|
||||
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.CenterCrop(IMAGE_SIZE),
|
||||
transforms.ToTensor()
|
||||
])
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
model = Bird_CNN(c_in=3, c_hidden=16, c_out=7)
|
||||
full_path = os.path.join(BUILD_PATH, "bird_cnn")
|
||||
model.load_state_dict(torch.load(full_path, map_location=torch.device(device)))
|
||||
model.to(device)
|
||||
model.eval()
|
||||
|
||||
img = cv2.imread("./testimages/two_crows.jpg")
|
||||
|
||||
if img is None:
|
||||
raise ValueError("Image not found or path is wrong")
|
||||
|
||||
ss = cv2.ximgproc.segmentation.createSelectiveSearchSegmentation()
|
||||
ss.setBaseImage(img)
|
||||
|
||||
ss.switchToSelectiveSearchFast()
|
||||
rects = ss.process()
|
||||
|
||||
# convert to array for easy sorting
|
||||
rects = np.array(rects)
|
||||
|
||||
# compute area
|
||||
areas = rects[:, 2] * rects[:, 3]
|
||||
|
||||
# sort by area (descending)
|
||||
idx = np.argsort(-areas)
|
||||
|
||||
# take top 10
|
||||
top10 = rects[idx[:200]]
|
||||
class_conf_sum = defaultdict(float)
|
||||
|
||||
for (x, y, w, h) in top10:
|
||||
#cv2.rectangle(img_copy, (x, y), (x + w, y + h), (0, 255, 0), 1)
|
||||
crop = img[y:y+h, x:x+w]
|
||||
crop_pil = Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
|
||||
image = transform(crop_pil).unsqueeze(0).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
pred = model(image)
|
||||
probs = F.softmax(pred, dim=1)
|
||||
confidence, cls = torch.max(probs, dim=1)
|
||||
if confidence.item() < 0.7:
|
||||
continue
|
||||
|
||||
class_conf_sum[cls.item()] += confidence.item()
|
||||
|
||||
for i in range(6):
|
||||
print(str(i) + ": " + str(class_conf_sum[i]))
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 65 KiB |
Reference in New Issue
Block a user