warhawk_base.py

from ultralytics import YOLO as warhawkAi
import cvzone as warhawkAiInterpretor
import cv2 as warhawkAiMedia
import math as warhawkMath
import warhawkObjectClassnames_fire
import numpy as np
from sort import *

def log_unique_detections_to_file(filename, detections):
    logged_entries = set()

    # Read existing logs into the set
    try:
        with open(filename, 'r') as file:
            for line in file:
                if line.strip():  # Ensure line is not empty
                    parts = line.split(', ')
                    if len(parts) == 2:
                        id_str, coords_str = parts
                        id = int(id_str.split(': ')[1])
                        coords = tuple(map(int, coords_str.split(': ')[1][1:-1].split(', ')))
                        logged_entries.add((id, coords))
    except FileNotFoundError:
        pass  # Handle if the file doesn't exist initially

    # Append new unique detections to the file
    with open(filename, 'a') as file:
        for detection in detections:
            id, coords = detection
            if (id, coords) not in logged_entries:
                logged_entries.add((id, coords))
                file.write(f"ID: {id}, Coordinates: ({coords[0]}, {coords[1]}, {coords[2]}, {coords[3]})\n")


def detectVideo(path_x):
    videoPath = path_x
    warhawkAiClassnames = warhawkObjectClassnames_fire.warhawkAiClassnames
    droneFeedCapture = warhawkAiMedia.VideoCapture(videoPath)
    warhawkModel = warhawkAi("/home/sarin/PycharmProjects/pythonProject/WarHawk-AI/warhawk-weights/real_avatar.pt")
    warhawkAiClassesToDetect = [0, 1, 2, 3]

    humanTracker = Sort(max_age=40, min_hits=8, iou_threshold=0.1)
    weaponTracker = Sort(max_age=40, min_hits=6, iou_threshold=0.1)
    treesTracker = Sort(max_age=40, min_hits=6, iou_threshold=0.1)

    while True:
        success, img = droneFeedCapture.read()
        results = warhawkModel(img, stream=True, conf=0.2, classes=warhawkAiClassesToDetect)
        humanDetections = np.empty((0, 5))
        weaponDetections = np.empty((0, 5))
        treesDetections = np.empty((0, 5))

        hsvFrame = warhawkAiMedia.cvtColor(img, warhawkAiMedia.COLOR_BGR2HSV)

        # Tracking and displaying objects
        for r in results:
            for box in r.boxes:
                x1, y1, x2, y2 = box.xyxy[0]
                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
                conf = warhawkMath.ceil((box.conf[0] * 100)) / 100
                objectClass = warhawkAiClassnames[int(box.cls[0])]
                
                if objectClass == 'person':
                    # Set range for red color and define mask
                    red_lower = np.array([136, 87, 111], np.uint8)
                    red_upper = np.array([180, 255, 255], np.uint8)
                    red_mask = warhawkAiMedia.inRange(hsvFrame, red_lower, red_upper)
                    red_mask = warhawkAiMedia.dilate(red_mask, np.ones((5, 5), np.uint8))
                    res_red = warhawkAiMedia.bitwise_and(img, img, mask=red_mask)

                    # Set range for green color and define mask
                    green_lower = np.array([25, 52, 72], np.uint8)
                    green_upper = np.array([102, 255, 255], np.uint8)
                    green_mask = warhawkAiMedia.inRange(hsvFrame, green_lower, green_upper)
                    green_mask = warhawkAiMedia.dilate(green_mask, np.ones((5, 5), np.uint8))
                    res_green = warhawkAiMedia.bitwise_and(img, img, mask=green_mask)

                    # Set range for blue color and define mask
                    blue_lower = np.array([94, 80, 2], np.uint8)
                    blue_upper = np.array([120, 255, 255], np.uint8)
                    blue_mask = warhawkAiMedia.inRange(hsvFrame, blue_lower, blue_upper)
                    blue_mask = warhawkAiMedia.dilate(blue_mask, np.ones((5, 5), np.uint8))
                    res_blue = warhawkAiMedia.bitwise_and(img, img, mask=blue_mask)

                    # Draw rectangles and labels for detections
                    warhawkAiInterpretor.cornerRect(img, (x1, y1, x2 - x1, y2 - y1), l=9, rt=2, t=2, colorR=(225, 100, 255))
                    warhawkAiInterpretor.putTextRect(img, f'{objectClass} {conf}', (max(0, x1), max(0, y1)), scale=1,
                                                     thickness=1, offset=10, colorR=(225, 100, 225),
                                                     font=warhawkAiMedia.FONT_HERSHEY_PLAIN)

                    # Display red objects inside the person's bounding box
                    contours, hierarchy = warhawkAiMedia.findContours(red_mask[y1:y2, x1:x2], warhawkAiMedia.RETR_TREE, warhawkAiMedia.CHAIN_APPROX_SIMPLE)
                    for pic, contour in enumerate(contours):
                        area = warhawkAiMedia.contourArea(contour)
                        if area > 10:
                            cx, cy, cw, ch = warhawkAiMedia.boundingRect(contour)
                            cx += x1  # Adjust contour coordinates relative to the image
                            cy += y1  # Adjust contour coordinates relative to the image
                            img = warhawkAiMedia.rectangle(img, (cx, cy), (cx + cw, cy + ch), (0, 0, 255), 2)

                    # Display green objects inside the person's bounding box
                    contours, hierarchy = warhawkAiMedia.findContours(green_mask[y1:y2, x1:x2], warhawkAiMedia.RETR_TREE, warhawkAiMedia.CHAIN_APPROX_SIMPLE)
                    for pic, contour in enumerate(contours):
                        area = warhawkAiMedia.contourArea(contour)
                        if area > 10:
                            cx, cy, cw, ch = warhawkAiMedia.boundingRect(contour)
                            cx += x1  # Adjust contour coordinates relative to the image
                            cy += y1  # Adjust contour coordinates relative to the image
                            img = warhawkAiMedia.rectangle(img, (cx, cy), (cx + cw, cy + ch), (0, 255, 0), 2)

                    # Display blue objects inside the person's bounding box
                    contours, hierarchy = warhawkAiMedia.findContours(blue_mask[y1:y2, x1:x2], warhawkAiMedia.RETR_TREE, warhawkAiMedia.CHAIN_APPROX_SIMPLE)
                    for pic, contour in enumerate(contours):
                        area = warhawkAiMedia.contourArea(contour)
                        if area > 10:
                            cx, cy, cw, ch = warhawkAiMedia.boundingRect(contour)
                            cx += x1  # Adjust contour coordinates relative to the image
                            cy += y1  # Adjust contour coordinates relative to the image
                            img = warhawkAiMedia.rectangle(img, (cx, cy), (cx + cw, cy + ch), (255, 0, 0), 2)

                    currentArray = np.array([x1, y1, x2, y2, conf])
                    humanDetections = np.vstack((humanDetections, currentArray))

                elif objectClass == 'stick':
                    warhawkAiInterpretor.cornerRect(img, (x1, y1, x2 - x1, y2 - y1), l=9, rt=2, t=2, colorR=(8, 33, 12))
                    warhawkAiInterpretor.putTextRect(img, f'Tree {conf}', (max(0, x1), max(35, y1)), scale=1,
                                                     thickness=1, offset=10, colorR=(8, 33, 12),
                                                     font=warhawkAiMedia.FONT_HERSHEY_PLAIN)
                    currentArray = np.array([x1, y1, x2, y2, conf])
                    treesDetections = np.vstack((treesDetections, currentArray))
                elif objectClass == 'tree':
                    warhawkAiInterpretor.cornerRect(img, (x1, y1, x2 - x1, y2 - y1), l=9, rt=2, t=2, colorR=(0, 0, 255))
                    warhawkAiInterpretor.putTextRect(img, f'Weapon {conf}', (max(0, x1), max(35, y1)), scale=1,
                                                     thickness=1, offset=10, colorR=(0, 0, 225),
                                                     font=warhawkAiMedia.FONT_HERSHEY_PLAIN)
                    currentArray = np.array([x1, y1, x2, y2, conf])
                    weaponDetections = np.vstack((weaponDetections, currentArray))

        humanOutputTracker = humanTracker.update(humanDetections)
        weaponOutputTracker = weaponTracker.update(weaponDetections)
        treesOutputTracker = treesTracker.update(treesDetections)

        # Log the object positions
        humanLogs = [(int(detection[4]), (int(detection[0]), int(detection[1]), int(detection[2]), int(detection[3]))) for detection in humanOutputTracker]
        weaponLogs = [(int(detection[4]), (int(detection[0]), int(detection[1]), int(detection[2]), int(detection[3]))) for detection in weaponOutputTracker]
        treesLogs = [(int(detection[4]), (int(detection[0]), int(detection[1]), int(detection[2]), int(detection[3]))) for detection in treesOutputTracker]

        log_unique_detections_to_file('human_logs.txt', humanLogs)
        log_unique_detections_to_file('weapon_logs.txt', weaponLogs)
        log_unique_detections_to_file('trees_logs.txt', treesLogs)

        # Display the frame
        warhawkAiMedia.imshow("WarHawk AI", img)
        if warhawkAiMedia.waitKey(1) & 0xFF == ord('q'):
            break

    droneFeedCapture.release()
    warhawkAiMedia.destroyAllWindows()

# Example usage
detectVideo("/home/sarin/PycharmProjects/pythonProject/WarHawk-AI/media/videos/Video_2024-07-18_15-27-57.mp4")