86 lines
1.9 KiB
Python
86 lines
1.9 KiB
Python
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]))
|