display.py

import cv2 as cv
import tensorflow as tf
import uuid
import datetime
import os
import numpy as np
# from keras_cv import bounding_box, visualization

class_mapping_smoker = {
    0: 'person',
    1: 'smoke'
}
color_mapping_smoker = {
    'person': (0, 255, 0),
    'smoke': (0, 0, 255)
}

def load_image(image_path):
    image = tf.io.read_file(image_path)
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.image.resize(image, (640,640))
    return image

def draw_bouding_box(frame,  pred):
    is_save = False
    selected_indices = [0]
    if pred['num_detections']:
        selected_indices = tf.image.non_max_suppression(
            list(pred['boxes']), list(pred['confidence']), 8,iou_threshold = .45)
        print(selected_indices)
    for i in range(pred['num_detections']):
        if not i in list(selected_indices):
            continue
        confidence = pred['confidence'][i].numpy()
        class_name = class_mapping_smoker[pred['classes'][i].numpy()]
        if class_name == 'person':
            if confidence < 0.45:
                continue
        else:
            if confidence < 0.25:
                continue
            else:
                is_save = True

        box = tf.cast(pred['boxes'][i], tf.int32).numpy()
        x1, y1, x2, y2 = box
        color = color_mapping_smoker[class_name]  # Green color
        thickness = 2  # Line thickness
        cv.rectangle(frame, (x1, y1), (x2, y2), color, thickness)
        text = f"Conf :{confidence:.2f} | {class_name}"
        cv.putText(frame, text, (x1, y1 - 10), cv.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    if is_save:
       save_frame(frame)


    return frame




def save_frame(frame):
    folder_name = datetime.date.today().isoformat()
    folder_path = os.path.join(os.getcwd(), 'images', folder_name)
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
    file_name = uuid.uuid4()
    cv.imwrite(os.path.join(folder_path, f'{file_name}.jpg'), frame)




def add_padding(frame, window_width=1080, window_height= 720):
    aspect_ratio = frame.shape[1] / frame.shape[0]

    # Calculate the scaling factors to fit the image within the desired dimensions
    scale_width = window_width / frame.shape[1]
    scale_height = window_height / frame.shape[0]

    # Choose the minimum scaling factor to ensure the entire image fits within the window
    scale_factor = min(scale_width, scale_height)

    # Resize the frame using the calculated scale factor
    new_width = int(frame.shape[1] * scale_factor)
    new_height = int(frame.shape[0] * scale_factor)
    frame = cv.resize(frame, (new_width, new_height))

    # Create a black background with the desired dimensions
    background = 255 * np.ones((window_height, window_width, 3), dtype=np.uint8)

    # Calculate the position to place the resized frame in the center of the window
    x_offset = (window_width - new_width) // 2
    y_offset = (window_height - new_height) // 2

    # Place the resized frame on the black background at the calculated position
    background[y_offset:y_offset + new_height, x_offset:x_offset + new_width] = frame

    return background





# def draw_bouding_box(frame,  pred):
#     is_save = False
#     for i in range(pred['num_detections']):
#         box = tf.cast(pred['boxes'][i], tf.int32).numpy()
#         confidence = pred['confidence'][i].numpy()
#         if confidence < .4:
#             continue
#         class_name = class_mapping_smoker[pred['classes'][i].numpy()]
#         if class_name == 'person':
#             if confidence < 0.5:
#                 continue
#         else:
#             if confidence < 0.3:
#                 continue
#             else:
#                 is_save = True
#         x1, y1, x2, y2 = box
#         color = color_mapping_smoker[class_name]  # Green color
#         thickness = 2  # Line thickness
#         cv.rectangle(frame, (x1, y1), (x2, y2), color, thickness)
#         text = f"Conf :{confidence:.2f} | {class_name}"
#         cv.putText(frame, text, (x1, y1 - 10), cv.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
#
#     if is_save:
#        save_frame(frame)
#
#
#     return frame

#
# # def visualize_detections(model, images, bounding_box_format):
# #     y_pred = model.predict(images)
# #     y_pred = {'boxes': tf.ragged.constant([y_pred['boxes'][y_pred['confidence'] != -1]]),
# #               'confidence': tf.ragged.constant([y_pred['confidence'][y_pred['confidence'] != -1]]),
# #               'classes': tf.ragged.constant([y_pred['classes'][y_pred['confidence'] != -1]]),
# #               'num_detections': np.array([np.count_nonzero(y_pred['confidence'] != -1)], dtype=np.int32)
# #               }
# #     print(y_pred)
# #     visualization.plot_bounding_box_gallery(
# #         images,
# #         value_range=(0, 255),
# #         bounding_box_format=bounding_box_format,
# #         # y_true=y_true,
# #         y_pred=y_pred,
# #         scale=4,
# #         rows=1,
# #         cols=1,
# #         show=True,
# #         font_scale=0.7,
# #         class_mapping=class_mapping,
# #     )