import cv2 import numpy as np import time import json import os import random import torch from torchvision import transforms import torch.nn as nn with open("config.json", "r") as f: CONFIG = json.load(f) DEVICE = CONFIG["model"]["device"] if DEVICE == "auto": DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # a simple custom model class CustomCNN(nn.Module): def __init__(self): super(CustomCNN, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 1, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2) ) # self.classifier = nn.Sequential( nn.Linear(1 * 56 * 56, 32), # Adjust input features based on actual output size nn.ReLU(), nn.Linear(32, 2) # Assuming 2 classes ) def forward(self, x): x = self.features(x) x = torch.flatten(x, 1) # Flatten the tensor for the fully connected layer x = self.classifier(x) return x model = CustomCNN().to(DEVICE) model.load_state_dict(torch.load(CONFIG["model"]["path"], map_location=torch.device(DEVICE) ) ) quant_dtype = CONFIG["model"]["quantization"] if quant_dtype == "qint8": model = torch.quantization.quantize_dynamic( model, None, dtype=torch.qint8 ) model.eval() print("model loaded") class BirdDetection(): def __init__(self, threshold=CONFIG["bird_detection"]["threshold"], percentile=CONFIG["bird_detection"]["percentile"], classification_model=model, horizon_height_ratio=CONFIG["bird_detection"]["horizon_height_ratio"], min_bbox_dim=CONFIG["bird_detection"]["min_bbox_dim"], max_bboxes = CONFIG["bird_detection"]["max_bboxes"], ): self.threshold = threshold self.percentile = percentile self.transform = transforms.Compose([ transforms.ToPILImage(), transforms.Resize((224, 224)), # Resize images to a fixed size transforms.ToTensor(), # Convert images to PyTorch tensors ]) self.classification_model = classification_model self.change_inf_time = [] self.contour_inf_time = [] self.class_inf_time = [] self.horizon_height_ratio = horizon_height_ratio self.min_bbox_dim = min_bbox_dim self.max_bboxes = max_bboxes def binary_image_to_coco_bboxes(self, binary_img): """ Args: binary_img (np.ndarray): Binary image (2D array with values 0 and 255 or 0 and 1) Returns: List[Dict]: A list of dictionaries each with keys: {'category_id': 0, 'bbox': [x_min, y_min, width, height]} """ # Ensure binary image is in the correct format for contour detection img_height = binary_img.shape[0] horizon_height = int(img_height * self.horizon_height_ratio) binary_img[horizon_height:, :] = 0 binary_img = (binary_img * 255).astype(np.uint8) contours, _ = cv2.findContours(binary_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE ) coco_bboxes = [] for contour in contours: x, y, w, h = cv2.boundingRect(contour) if w < self.min_bbox_dim: x = x - ((self.min_bbox_dim-w)//2) w = self.min_bbox_dim if h < self.min_bbox_dim: y = y - ((self.min_bbox_dim-h)//2) h = self.min_bbox_dim coco_bboxes.append({ 'category_id': 0, 'bbox': [float(x), float(y), float(w), float(h)] }) return coco_bboxes def classify_image_bboxes(self,img, coco_bboxes): if len(coco_bboxes) == 0: return [] cropped_list = [] for bbox in coco_bboxes: x, y, w, h = bbox['bbox'] cropped_img = img[int(y):int(y+h), int(x):int(x+w)] cropped_list.append( self.transform(cropped_img) ) batch = torch.stack(cropped_list).to(DEVICE) outputs = self.classification_model(batch) _, predicted = torch.max(outputs.data, 1) correct_bboxes = [] for i, pred in enumerate(predicted): if pred == 1: correct_bboxes.append(coco_bboxes[i]) return correct_bboxes def __call__(self, img,prev_img,torelable_min = None, torelable_max = None): t1 = time.time() img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) prev_img_gray = cv2.cvtColor(prev_img, cv2.COLOR_BGR2GRAY) diff = cv2.absdiff(img_gray, prev_img_gray) diff_min, diff_max = diff.min(), diff.max() diff = (diff - diff_min) / (diff_max - diff_min) diff_mean = diff.mean() self.change_inf_time.append(time.time()-t1) if True: t1 = time.time() mask_threshold = np.percentile(diff[diff >= self.threshold], self.percentile) coco_bboxes = self.binary_image_to_coco_bboxes(diff >= mask_threshold) total_n_boxes = len(coco_bboxes) self.contour_inf_time.append((total_n_boxes,time.time()-t1)) t1 = time.time() coco_bboxes_raw = coco_bboxes.copy() if len(coco_bboxes_raw)>self.max_bboxes: coco_bboxes_raw = random.choices(coco_bboxes_raw,k=self.max_bboxes) coco_bboxes = self.classify_image_bboxes(img, coco_bboxes_raw) #TODO coco_bboxes += self.classify_image_bboxes(prev_img, coco_bboxes_raw) self.class_inf_time.append((total_n_boxes,time.time()-t1)) return diff, coco_bboxes def draw_coco_bboxes(image, annotations, color=(255, 0, 0), thickness=50): """ Args: image (np.ndarray): The input image (BGR format). annotations (List[Dict]): List of annotations with COCO bboxes: [{'category_id': int, 'bbox': [x, y, w, h]}, ...] color (Tuple[int, int, int]): BGR color for the boxes (default is red). thickness (int): Line thickness of the boxes. Returns: np.ndarray: The image with bounding boxes drawn on it. """ img_copy = image.copy() for ann in annotations: x, y, w, h = ann['bbox'] top_left = (int(round(x)), int(round(y))) bottom_right = (int(round(x + w)), int(round(y + h))) cv2.rectangle(img_copy, top_left, bottom_right, color, thickness) return img_copy bird_detector = BirdDetection() class Agent: def __init__(self,sensor, min_angle=CONFIG["agent"]["min_angle"], max_angle=CONFIG["agent"]["max_angle"], stride=CONFIG["agent"]["stride"], yaws = [0,0], bird_detector=None ): self.min_angle = min_angle self.max_angle = max_angle self.stride = stride self.yaws = yaws if bird_detector is None: self.bird_detector = BirdDetection() else: self.bird_detector = bird_detector self.sensor = sensor def extract_face(self, yaw=0.0): face = self.sensor.take_photo() return face def predict(self,current_image, prev_image): diff, coco_bboxes = self.bird_detector(current_image,prev_image) return coco_bboxes def handle_prediction(self,current_image, prev_image, coco_bboxes, yaw): viz_config = CONFIG["agent"]["visualization"] image = draw_coco_bboxes( current_image, coco_bboxes, tuple(viz_config["bbox_color"]), viz_config["bbox_thickness"] ) preview_size = tuple(viz_config["preview_size"]) return cv2.resize(image, preview_size, interpolation=cv2.INTER_NEAREST) def work(self): turn = 1 while True: yaws = self.yaws[::turn] turn = -turn for yaw in (yaws[1:]): prev_image = self.extract_face(yaw) current_image = self.extract_face(yaw) coco_bboxes = self.predict(current_image, prev_image) artifacts = self.handle_prediction(current_image, prev_image, coco_bboxes,yaw) yield artifacts