agent.py

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