frame_process.py

from queue import Queue
import cv2
import time
from detection_tracking import DetectionTracking
import os
import numpy as np
from PIL import Image
from reid import extract_features_from_pil, reidentification
import json


def process(frame, camera_idx, model, roi, detector, frame_counter, fps, track_id_duration_map, track_id_map, save_path):
    """
        Processes a frame for object detection and tracking.

        This function takes a frame and performs object detection and tracking on it. It also handles the saving of cropped images
        and extraction of features from the detected objects.

        Args:
            frame: The frame to be processed.
            camera_idx: The index of the camera source.
            model: The model used for re-identification.
            roi: The region of interest in the frame.
            detector: The detector used for object detection and tracking.
            frame_counter: The current frame count.
            fps: The frames per second of the video source.
            track_id_duration_map: A map that stores the duration of each track.
            track_id_map: A map that stores the re-identified track IDs.
            save_path: The path where the processed frames are saved.

        Returns:
            The processed frame with bounding boxes and IDs drawn on the detected objects.
        """
    x, y, w, h = roi
    frame_for_cropping = frame.copy()

    tracks = detector.process_frame(frame)
    for track in tracks:
        bbox = track.to_tlbr()  # Use the to_tlbr method to get the bounding box
        track_id = track.track_id  # Access the track_id property directly

        # Clamp the bounding box coordinates
        x_min, y_min, x_max, y_max = bbox
        x_min = max(0, min(x_min, frame.shape[1] - 1))
        y_min = max(0, min(y_min, frame.shape[0] - 1))
        x_max = max(0, min(x_max, frame.shape[1] - 1))
        y_max = max(0, min(y_max, frame.shape[0] - 1))

        adjusted_bbox = (x_min, y_min, x_max, y_max)
        center_x, center_y = ((x_min + x_max) / 2, (y_min + y_max) / 2)
        # adjusted_bbox = (bbox[0], bbox[1], bbox[2], bbox[3])
        # center_x, center_y = (adjusted_bbox[0] + adjusted_bbox[2]) / 2, (adjusted_bbox[1] + adjusted_bbox[3]) /2

        if frame_counter % 5 == 0:
            if x <= center_x <= x + w and y <= center_y <= y + h:
                cropped_img = frame_for_cropping[int(adjusted_bbox[1]):int(adjusted_bbox[3]),
                              int(adjusted_bbox[0]):int(adjusted_bbox[2])]
                if cropped_img.size<=0:
                    continue

                if camera_idx == 0:
                    gallery_folder = "Gallery/main_gallery"
                    os.makedirs(gallery_folder, exist_ok=True)
                    # if track_id not in track_id_map:
                    #     reidentified_id = reidentification(cropped_img, model, gallery_folder='Gallery/main_gallery', track_id=track_id)
                    #     track_id_map[track_id] = reidentified_id
                    track_id_map[track_id] = track_id
                    filename = f"{track_id_map[track_id]}_{camera_idx}_{frame_counter}.jpg"
                    filepath = os.path.join(gallery_folder, filename)
                    cv2.imwrite(filepath, cropped_img)

                    feature = extract_features_from_pil(Image.fromarray(cropped_img), model)
                    feature_path = os.path.join(save_path, "gallery_features", filename.replace('.jpg', '.npy'))
                    os.makedirs(os.path.dirname(feature_path), exist_ok=True)
                    np.save(feature_path, feature)
                else:
                    if track_id not in track_id_map:
                        reidentified_id = reidentification(cropped_img, model, gallery_folder='Gallery/main_gallery', track_id=track_id)
                        track_id_map[track_id] = reidentified_id
                    gallery_folder = os.path.join(save_path, "temporary_gallery")
                    os.makedirs(gallery_folder, exist_ok=True)
                    filename = f"{track_id_map[track_id]}_{camera_idx}_{frame_counter}.jpg"
                    filepath = os.path.join(gallery_folder, filename)
                    cv2.imwrite(filepath, cropped_img)

        print(f"Track ID: {track_id}, BBox: {adjusted_bbox}")

        if track_id not in track_id_duration_map:
            track_id_duration_map[track_id] = {'start': frame_counter, 'end': frame_counter}
        else:
            # Update the end frame counter for existing tracks
            track_id_duration_map[track_id]['end'] = frame_counter

            # Calculate duration for existing track IDs
        duration_frames = track_id_duration_map[track_id]['end']-track_id_duration_map[track_id]['start']
        duration_seconds = duration_frames / fps

        if track_id in track_id_map:
            re_track_id = track_id_map[track_id]
        else:
            re_track_id = track_id

        if track_id is not None:
            cv2.rectangle(frame, (int(adjusted_bbox[0]), int(adjusted_bbox[1])),(int(adjusted_bbox[2]), int(adjusted_bbox[3])), (0, 255, 0), 2)
            duration_text = f"ID: {re_track_id}, Time: {duration_seconds:.2f}s"
            cv2.putText(frame, duration_text, (int(bbox[0]), int(bbox[1]) - 10),cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)

    return frame


def generate_reid_report_files(track_id_maps, track_id_duration_maps, fps):
    """
        Generates report files for re-identified objects.

        This function takes the maps of track IDs and durations and generates a report file for each re-identified object.

        Args:
            track_id_maps: A list of maps that store the re-identified track IDs for each camera source.
            track_id_duration_maps: A list of maps that store the duration of each track for each camera source.
            fps: The frames per second of the video sources.
        """
    os.makedirs('log', exist_ok=True)  # Create a log folder

    reid_info = {}
    # Consolidate information for each Re_ID across all camera sources
    for source_idx, (track_map, duration_map) in enumerate(zip(track_id_maps, track_id_duration_maps)):
        for track_id, reid_id in track_map.items():
            if reid_id not in reid_info:
                reid_info[reid_id] = []

            if track_id in duration_map:
                start_frame = duration_map[track_id]['start']
                end_frame = duration_map[track_id]['end']
                in_time = start_frame / fps
                out_time = end_frame / fps
                reid_info[reid_id].append((source_idx, track_id, in_time, out_time))

    # Write to text files
    for reid_id, infos in reid_info.items():
        with open(f'log/{reid_id}.txt', 'w') as file:
            file.write(f"Re_ID = {reid_id}\n")
            for info in infos:
                file.write(
                    f"Camera_source idx: {info[0]}, Track_id: {info[1]}, In: {info[2]:.2f}s - Out: {info[3]:.2f}s\n")




def save_data_to_json(data, file_name):
    """
        Saves data to a JSON file.

        This function takes a data object and a file name, and writes the data to the file in JSON format.

        Args:
            data: The data to be saved.
            file_name: The name of the file where the data is saved.
        """
    with open(file_name, 'w') as file:
        json.dump(data, file, indent=4)