import sqlite3
import datetime
from Status import Status
import numpy as np

class Database:
    """
    A class to handle interactions with the SQLite database for object tracking.

    Attributes:
        conn (sqlite3.Connection): The database connection.
        cursor (sqlite3.Cursor): The database cursor.
    """
    def __init__(self, db_path='object_tracking.db'):
        """
        Initializes the Database class with a connection to the specified SQLite database.

        Args:
            db_path (str): The path to the SQLite database file.
        """
        self.conn = sqlite3.connect(db_path)
        self.cursor = self.conn.cursor()
        self.create_table()
        self.add_person_name_column()
        self.create_known_faces_table()

    def create_table(self):
        """
        Creates the 'tracked_objects' table if it does not already exist.
        """
        self.cursor.execute('''
        CREATE TABLE IF NOT EXISTS tracked_objects (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            class_name TEXT,
            x1 INTEGER,
            y1 INTEGER,
            width INTEGER,
            height INTEGER,
            last_seen DATETIME,
            status TEXT,
            detection_count INTEGER DEFAULT 0,
            person_name TEXT  -- Added new field for person's name
        )
    ''')
        self.conn.commit()
        
    def create_known_faces_table(self):
        """
        Creates the 'known_faces' table if it does not already exist.
        """
        self.cursor.execute('''
            CREATE TABLE IF NOT EXISTS known_faces (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                embedding BLOB NOT NULL
            )
        ''')
        self.conn.commit()

    def fetch_objects(self, class_name=None):
        """
        Fetches objects from the database, optionally filtered by class name.

        Args:
            class_name (str, optional): The class name to filter objects.

        Returns:
            list: A list of tuples representing the objects.
        """
        if class_name:
            self.cursor.execute('SELECT * FROM tracked_objects WHERE class_name=?', (class_name,))
        else:
            self.cursor.execute('SELECT * FROM tracked_objects')
        return self.cursor.fetchall()

    def fetch_all_objects(self):
        """
        Fetches all objects from the database.

        Returns:
            list: A list of tuples representing all objects.
        """
        self.cursor.execute('SELECT * FROM tracked_objects')
        return self.cursor.fetchall()

    def fetch_object_by_id(self, obj_id):
        """
        Fetches a single object by its ID.

        Args:
            obj_id (int): The ID of the object.

        Returns:
            tuple: The object data as a tuple.
        """
        self.cursor.execute('SELECT * FROM tracked_objects WHERE id=?', (obj_id,))
        return self.cursor.fetchone()

    def update_object(self, obj_id, x1, y1, width, height):
        """
        Updates the position and size of an existing object in the database, and increments detection_count.

        Args:
            obj_id (int): The ID of the object to update.
            x1 (int): The new x-coordinate of the object's top-left corner.
            y1 (int): The new y-coordinate of the object's top-left corner.
            width (int): The new width of the object.
            height (int): The new height of the object.

        Returns:
            int: The ID of the updated object.
        """
        last_seen = datetime.datetime.now().isoformat()
        status = Status.IN_FRAME.value  # Set status to 'In Frame'
        self.cursor.execute('''
            UPDATE tracked_objects 
            SET x1=?, y1=?, width=?, height=?, last_seen=?, status=?, detection_count=detection_count+1 
            WHERE id=?
        ''', (x1, y1, width, height, last_seen, status, obj_id))
        self.conn.commit()
        return obj_id

    def insert_object(self, class_name, x1, y1, width, height, person_name=None):
        """
        Inserts a new object into the database, and returns the new object's ID.

        Args:
            class_name (str): The class name of the object.
            x1 (int): The x-coordinate of the object's top-left corner.
            y1 (int): The y-coordinate of the object's top-left corner.
            width (int): The width of the object.
            height (int): The height of the object.
            person_name (str, optional): The name of the person associated with the object.

        Returns:
            int: The ID of the newly inserted object.
        """
        last_seen = datetime.datetime.now().isoformat()
        status = Status.IN_FRAME.value
        detection_count = 1
        self.cursor.execute('''
            INSERT INTO tracked_objects (class_name, x1, y1, width, height, last_seen, status, detection_count, person_name) 
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (class_name, x1, y1, width, height, last_seen, status, detection_count, person_name))
        self.conn.commit()
        return self.cursor.lastrowid

    def update_status(self, obj_id, status):
        """
        Updates the status of an object in the database.

        Args:
            obj_id (int): The ID of the object to update.
            status (Status): The new status of the object.
        """
        self.cursor.execute('''
            UPDATE tracked_objects 
            SET status=? 
            WHERE id=?
        ''', (status.value, obj_id))
        self.conn.commit()

    def delete_object(self, obj_id):
        """
        Deletes an object from the database.

        Args:
            obj_id (int): The ID of the object to delete.
        """
        self.cursor.execute('''
            DELETE FROM tracked_objects WHERE id=?
        ''', (obj_id,))
        self.conn.commit()

    def add_person_name_column(self):
        """
        Adds the 'person_name' column to 'tracked_objects' table if it doesn't exist.
        """
        self.cursor.execute("PRAGMA table_info(tracked_objects)")
        columns = [column[1] for column in self.cursor.fetchall()]
        if 'person_name' not in columns:
            self.cursor.execute('ALTER TABLE tracked_objects ADD COLUMN person_name TEXT')
            self.conn.commit()
            
    def insert_known_face(self, name, embedding):
        """
        Inserts a new known face into the database.

        Args:
            name (str): The name of the person.
            embedding (np.ndarray): The face embedding.
        """
        embedding_blob = embedding.astype(np.float32).tobytes()
        self.cursor.execute('''
            INSERT INTO known_faces (name, embedding) VALUES (?, ?)
        ''', (name, embedding_blob))
        self.conn.commit()

    def fetch_known_faces(self):
        """
        Fetches all known faces from the database.

        Returns:
            list: A list of tuples representing known faces.
        """
        self.cursor.execute('SELECT * FROM known_faces')
        known_faces = []
        for row in self.cursor.fetchall():
            id = row[0]
            name = row[1]
            embedding_blob = row[2]
            embedding = np.frombuffer(embedding_blob, dtype=np.float32)
            known_faces.append((id, name, embedding))
        return known_faces
    
    def get_known_face_by_name(self, name):
        """
        Fetches a known face by name from the database.

        Args:
            name (str): The name of the person.

        Returns:
            tuple: A tuple representing the known face.
        """
        self.cursor.execute('SELECT * FROM known_faces WHERE name=?', (name,))
        row = self.cursor.fetchone()
        if row:
            id = row[0]
            name = row[1]
            embedding_blob = row[2]
            embedding = np.frombuffer(embedding_blob, dtype=np.float32)
            return (id, name, embedding)
        return None
    
    def update_person_name(self, obj_id, person_name):
        self.cursor.execute('''
            UPDATE tracked_objects 
            SET person_name=? 
            WHERE id=?
        ''', (person_name, obj_id))
        self.conn.commit() import cv2
import numpy as np
from insightface.app import FaceAnalysis

class FaceInterface:
    """
    A class for interfacing with the face recognition module.
    """
    def __init__(self, db) -> None:
        """
        Initializes the FaceInterface with the given database.

        Args:
            db (Database): The database instance.
        """
        self.db = db
        # Initialize FaceAnalysis for face recognition
        self.face_recognition = FaceAnalysis(allowed_modules=['detection', 'recognition'])
        self.face_recognition.prepare(ctx_id=0, det_size=(320, 320))

    def add_known_face(self, image_path: str, name: str) -> bool:
        """
        Adds a known face to the database.

        Args:
            image_path (str): Path to the image file containing the face.
            name (str): Name of the person.

        Returns:
            bool: True if the face was successfully added, False otherwise.
        """
        # Read the image from the given path
        img = cv2.imread(image_path)
        if img is None:
            print(f"Failed to read image at {image_path}")
            return False

        # Detect faces and get embeddings
        faces = self.face_recognition.get(img)
        if not faces:
            print(f"No face detected in image {image_path}")
            return False

        # Use the first detected face
        face = faces[0]
        embedding = face.embedding
        if embedding is None:
            print(f"Failed to get embedding for face in image {image_path}")
            return False

        # Normalize the embedding
        embedding = embedding / np.linalg.norm(embedding)
        
        if self.db.get_known_face_by_name(name):
            print(f"Face with name {name} already exists in the database")
            return True

        # Insert the name and embedding into the database
        self.db.insert_known_face(name, embedding)
        return True. # Import necessary libraries
from Database import Database
from ObjectDetection import ObjectDetection
from ObjectTracker import ObjectTracker
from FaceInterface import FaceInterface
import cv2
import argparse

def add_known_faces(file_path: str, face_interface: FaceInterface) -> None:
    """
    Add known faces from the CSV file to the face recognition module.

    Args:
        file_path (str): Path to the CSV file containing image paths and names.
        face_interface (FaceInterface): The face recognition interface.
        
    Raises:
        ValueError: If a known face could not be added.
    """
    if file_path:
        with open(file_path, 'r') as file:
            for line in file:
                image_path, name = line.strip().split(',')
                if not face_interface.add_known_face(image_path, name): 
                    raise ValueError(f"Failed to add {name} from {image_path}")

def main():
    """
    The main function to initialize object detection and tracking,
    and process the video stream from the webcam.
    """
    # Set up argument parser
    parser = argparse.ArgumentParser(description="CSV file of image_path, name. For face recognition.")
    parser.add_argument('--csv', type=str, help='Path to the CSV file for face recognition.')
    parser.add_argument('--model', type=str, default='yolov8n.pt', help='The YOLO model weights name. Default is yolov8n.pt.')
    parser.add_argument('--show', action='store_true', help='Display the video stream.')
    args = parser.parse_args()
    
    # Initialize the database, detection model, and tracker
    db = Database()
    detection_model = ObjectDetection(model=args.model)
    tracker = ObjectTracker(db, detection_model)   
    face_interface = FaceInterface(db)
    
    # Add known faces from the CSV file
    add_known_faces(args.csv, face_interface)

    # Set up video capture from the default webcam
    cap = cv2.VideoCapture(0)
    cap.set(3, 640)  # Set width
    cap.set(4, 480)  # Set height

    while True:
        success, img = cap.read()
        if not success:
            break

        # Perform object tracking on the current frame
        img = tracker.track_objects(img)

        # Display the resulting frame
        if args.show:
            cv2.imshow('Webcam', img)

        # Exit if 'q' is pressed
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # Release resources
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main(). from ultralytics import YOLO
import os

class ObjectDetection:
    """
    A class for object detection using a YOLO model.

    Attributes:
        model (YOLO): The YOLO model used for detection.
        classNames (list): List of class names that the model can detect.
        allowed_classes (set): Set of class indices that are allowed for tracking.
    """
    def __init__(self, model="yolov8n.pt"):
        """
        Initializes the ObjectDetection class with the specified YOLO model.

        Args:
            model_path (str): The path to the YOLO model weights.
        """
        try:
            self.model = YOLO(self.get_ncnn_model(model), task='detect')
        except Exception as e:
            raise Exception(f"Error loading YOLO model: {e}")
        
        self.classNames = ["person", "bicycle", "car", "motorbike", "aeroplane", "bus", "train", "truck", "boat",
                           "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
                           "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
                           "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat",
                           "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
                           "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
                           "carrot", "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant", "bed",
                           "diningtable", "toilet", "tvmonitor", "laptop", "mouse", "remote", "keyboard", "cell phone",
                           "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors",
                           "teddy bear", "hair drier", "toothbrush"]
        # Only allow certain classes for tracking
        self.allowed_classes = {0, 24, 25, 26, 31, 36, 37, 39, 40, 41, 53, 63, 65, 67}
        
    def get_ncnn_model(self, model: str) -> str:
        """
        Converts the given YOLO model to NCNN format.

        Args:
            model (str): The name of the YOLO model file.

        Raises:
            ValueError: If the model conversion fails.

        Returns:
            str: The path to the converted ONNX model.
        """
        onnx_model_path = model.split('.')[0] + '_ncnn_model'
        if os.path.exists(onnx_model_path):
            return onnx_model_path
        try:
            model: YOLO = YOLO(model)
            return model.export(format='ncnn')
        except Exception as e:
            raise ValueError(f"Failed to convert model to NCNN format: {e}")

    def detect_objects(self, img):
        """
        Detects objects in the given image.

        Args:
            img (numpy.ndarray): The image in which to detect objects.

        Returns:
            list: A list of detections, each containing class name, bounding box coordinates, width, height, and confidence.
        """
        results = self.model(img, stream=True)
        detections = []
        for r in results:
            for box in r.boxes:
                # Get the confidence score
                confidence = box.conf[0].item()
                if confidence < 0.55:
                    continue  # Skip detections below confidence threshold

                # Get the class index
                cls = int(box.cls[0])
                if cls not in self.allowed_classes:
                    continue  # Skip classes that are not allowed

                # Get bounding box coordinates
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                width, height = x2 - x1, y2 - y1
                class_name = self.classNames[cls]
                detections.append((class_name, x1, y1, width, height, confidence))
        return detections. import cv2
from deep_sort_realtime.deepsort_tracker import DeepSort
from datetime import datetime, timedelta
from Status import Status
from insightface.app import FaceAnalysis
import numpy as np
import concurrent.futures


class ObjectTracker:
    """
    A class for tracking objects using Deep Sort and updating their status in the database.

    Attributes:
        db (Database): The database instance for storing object data.
        detection_model (ObjectDetection): The object detection model.
        tracker (DeepSort): The Deep Sort tracker instance.
        object_positions (dict): Dictionary to store the positions of each object.
    """
    def __init__(self, db, detection_model):
        """
        Initializes the ObjectTracker class with the given database and detection model.

        Args:
            db (Database): The database instance.
            detection_model (ObjectDetection): The object detection model.
        """
        self.db = db
        self.detection_model = detection_model
        self.tracker = DeepSort(max_age=10, n_init=3, nms_max_overlap=1.0,
                                max_cosine_distance=0.4, nn_budget=None)
        self.object_positions = {}  # To store positions of each object
        # Initialize ArcFace for face recognition
        self.face_recognizer = FaceAnalysis(allowed_modules=['detection', 'recognition'])
        self.face_recognizer.prepare(ctx_id=0, det_size=(320, 320))
        self.face_recognition_cache = {}  # Stores obj_id: (person_name, last_recognition_frame)
        self.face_recognition_interval = 10  # Number of frames to wait before re-running face recognition
        self.frame_count = 0  # Initialize frame counter
        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=2)

    def is_near(self, existing_obj, new_obj, threshold=30):
        """
        Determines if a new object is near an existing object based on a threshold.

        Args:
            existing_obj (tuple): The existing object's coordinates and size (x, y, w, h).
            new_obj (tuple): The new object's coordinates and size (x, y, w, h).
            threshold (int, optional): The distance threshold to consider.

        Returns:
            bool: True if the objects are near each other, False otherwise.
        """
        ex, ey, ew, eh = existing_obj
        nx, ny, nw, nh = new_obj
        # Check if bounding boxes are close enough
        return not (nx > ex + ew + threshold or
                    nx + nw < ex - threshold or
                    ny > ey + eh + threshold or
                    ny + nh < ey - threshold)

    def is_within_frame(self, obj_id, frame_width, frame_height, border_threshold=20, size_threshold=60):
        """
        Checks if an object is within the frame or near the borders. If the object is near a border and its size is
        small (less than size_threshold), it is considered out of bounds.

        Args:
            obj_id (int): The ID of the object.
            frame_width (int): The width of the frame.
            frame_height (int): The height of the frame.
            border_threshold (int, optional): The distance from borders to consider as 'near the border'.
            size_threshold (int, optional): The size threshold to consider the object as out of bounds if near border.

        Returns:
            bool: True if the object is within the frame, False if it's out of bounds.
        """
        positions = self.object_positions.get(obj_id, [])
        if not positions:
            return False  # No position data available

        x, y, w, h = positions[-1]  # Get the most recent position

        # Check distances from borders
        distance_left = x
        distance_right = frame_width - (x + w)
        distance_top = y
        distance_bottom = frame_height - (y + h)

        # If object is near left border
        if distance_left <= border_threshold and w <= size_threshold:
            return False  # Object is out of bounds
        # If object is near right border
        if distance_right <= border_threshold and w <= size_threshold:
            return False
        # If object is near top border
        if distance_top <= border_threshold and h <= size_threshold:
            return False
        # If object is near bottom border
        if distance_bottom <= border_threshold and h <= size_threshold:
            return False

        # If none of the above conditions met, object is within frame
        return True

    def check_for_false_positives(self, frame_width, frame_height):
        """
        Checks the database for objects that haven't been seen for more than 500 ms.
        If an object hasn't been seen for more than 500 ms and is not within frame,
        updates its status to 'Left Frame'. If detection_count < 5, deletes it.

        Args:
            frame_width (int): The width of the frame.
            frame_height (int): The height of the frame.
        """
        all_objects = self.db.fetch_all_objects()
        now = datetime.now()
        detection_timeout = now - timedelta(milliseconds=500)

        for obj in all_objects:
            obj_id = obj[0]
            class_name = obj[1]
            x1 = obj[2]
            y1 = obj[3]
            width = obj[4]
            height = obj[5]
            last_seen_str = obj[6]
            status = obj[7]
            detection_count = obj[8]

            last_seen = datetime.fromisoformat(last_seen_str)

            # If the object hasn't been seen for more than 500 ms
            if last_seen < detection_timeout:
                if detection_count < 6:
                    # Delete the object from the database
                    self.db.delete_object(obj_id)
                    # Remove from position tracking
                    if obj_id in self.object_positions:
                        del self.object_positions[obj_id]
                else:
                    # Check if the object is within the frame
                    if not self.is_within_frame(obj_id, frame_width, frame_height):
                        # Update the object's status to 'Left Frame'
                        self.db.update_status(obj_id, Status.LEFT_FRAME)
    
    def track_objects(self, img):
        self.frame_count += 1  # Update frame count
        frame_height, frame_width = img.shape[:2]
        detections = self.detection_model.detect_objects(img)
        formatted_detections = []

        # Fetch all existing objects once per frame
        existing_objs = self.db.fetch_all_objects()
        existing_objs_dict = {}  # Map class_name to list of objects
        for obj in existing_objs:
            class_name = obj[1]
            if class_name not in existing_objs_dict:
                existing_objs_dict[class_name] = []
            existing_objs_dict[class_name].append(obj)

        for class_name, x1, y1, width, height, confidence in detections:
            new_obj = (x1, y1, width, height)
            formatted_detections.append(([x1, y1, width, height], confidence, class_name))
            cv2.rectangle(img, (x1, y1), (x1 + width, y1 + height), (255, 0, 255), 3)
            cv2.putText(img, f'{class_name} {confidence:.2f}', (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)

            person_name = None  # Initialize person_name

            if class_name == 'person':
                objs_in_class = existing_objs_dict.get(class_name, [])
                preexisting_obj = False
                obj_id = None

                for existing_obj in objs_in_class:
                    db_x1 = existing_obj[2]
                    db_y1 = existing_obj[3]
                    db_width = existing_obj[4]
                    db_height = existing_obj[5]
                    last_seen = datetime.fromisoformat(existing_obj[6])
                    status = Status(existing_obj[7])

                    if self.is_near((db_x1, db_y1, db_width, db_height), new_obj):
                        preexisting_obj = True
                        obj_id = existing_obj[0]
                        self.db.update_object(obj_id, x1, y1, width, height)
                        self.object_positions[obj_id] = [(x1, y1, width, height)]
                        break

                if not preexisting_obj:
                    obj_id = self.db.insert_object(class_name, x1, y1, width, height)
                    self.object_positions[obj_id] = [(x1, y1, width, height)]
                    self.db.update_status(obj_id, Status.IN_FRAME)
                    # Perform face recognition for new object
                    self.executor.submit(self.async_face_recognition, obj_id, img, x1, y1, width, height)
                    if person_name:
                        self.db.update_person_name(obj_id, person_name)
                        self.face_recognition_cache[obj_id] = (person_name, self.frame_count)
                else:
                    last_recognition_frame = self.face_recognition_cache.get(obj_id, (None, -self.face_recognition_interval))[1]
                    if self.frame_count - last_recognition_frame >= self.face_recognition_interval:
                        person_name = self.perform_face_recognition(img, x1, y1, width, height)
                        if person_name:
                            self.db.update_person_name(obj_id, person_name)
                            self.face_recognition_cache[obj_id] = (person_name, self.frame_count)
                    else:
                        # Use cached name
                        person_name = self.face_recognition_cache.get(obj_id, (None, None))[0]

                if person_name and person_name != 'Unknown':
                    cv2.putText(img, f'Name: {person_name}', (x1, y1 - 30),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
            existing_objs = self.db.fetch_objects(class_name)
            preexisting_obj = False
            obj_id = None

            for existing_obj in existing_objs:
                db_x1 = existing_obj[2]
                db_y1 = existing_obj[3]
                db_width = existing_obj[4]
                db_height = existing_obj[5]
                last_seen = datetime.fromisoformat(existing_obj[6])
                status = Status(existing_obj[7])

                if self.is_near((db_x1, db_y1, db_width, db_height), new_obj):
                    preexisting_obj = True
                    obj_id = existing_obj[0]
                    self.db.update_object(obj_id, x1, y1, width, height)
                    if class_name == 'person':
                        self.db.update_person_name(obj_id, person_name)
                    self.object_positions[obj_id] = [(x1, y1, width, height)]
                    break

            if not preexisting_obj:
                obj_id = self.db.insert_object(class_name, x1, y1, width, height, person_name)
                self.object_positions[obj_id] = [(x1, y1, width, height)]
            else:
                self.object_positions[obj_id] = [(x1, y1, width, height)]
                self.db.update_status(obj_id, Status.IN_FRAME)

        tracks = self.tracker.update_tracks(formatted_detections, frame=img)
        self.check_for_false_positives(frame_width, frame_height)
        return img

    def perform_face_recognition(self, img, x1, y1, width, height):
        person_img = img[y1:y1+height, x1:x1+width]
        faces = self.face_recognizer.get(person_img)
        if faces:
            face = faces[0]
            face_embedding = face.embedding / np.linalg.norm(face.embedding)
            known_faces = self.db.fetch_known_faces()
            min_distance = float('inf')
            matched_name = None

            for known_face in known_faces:
                known_name = known_face[1]
                known_embedding = known_face[2] / np.linalg.norm(known_face[2])
                distance = np.linalg.norm(face_embedding - known_embedding)
                if distance < min_distance:
                    min_distance = distance
                    matched_name = known_name

            threshold = 1.0  # Adjust based on your model and data
            if min_distance < threshold:
                return matched_name
            else:
                return 'Unknown'
        else:
            return 'No Face Detected'
        
    def async_face_recognition(self, obj_id, img, x1, y1, width, height):
        person_name = self.perform_face_recognition(img, x1, y1, width, height)
        if person_name:
            self.db.update_person_name(obj_id, person_name)
            self.face_recognition_cache[obj_id] = (person_name, self.frame_count). from enum import Enum

class Status(str, Enum):
    """
    An enumeration to represent the status of tracked objects.
    """
    IN_FRAME = "In Frame"
    LEFT_FRAME = "Left Frame". . Modify this code so that instead of the detection threshold being 0.55 make it be less if the object is far away (has a small box).