changed website structure, added new cnn files

bird_cnn/selective_search.py

@@ -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]))