Person Segmentation

process.py

from itertools import zip_longest
import json
import math
from typing import Any
import cv2
from loguru import logger
from shapely.geometry import Polygon
import numpy as np

DEBUGGING = True
COLORS = [
    (255, 0, 0),  # Red
    (0, 255, 0),  # Green
    (0, 0, 255),  # Blue
    (255, 255, 0),  # Yellow
    (0, 255, 255),  # Cyan
    (255, 0, 255),  # Magenta
    (255, 165, 0),  # Orange
    (128, 0, 128),  # Purple
    (128, 128, 0),  # Olive
    (0, 128, 128),  # Teal
]

if DEBUGGING:
    cv2.namedWindow("main", cv2.WINDOW_AUTOSIZE)


def calculate_overlap(
    face_bb_group: list[Polygon],
    output: tuple[int, int],
) -> tuple[float, Polygon | None]:
    """
    Calculate overlap between two rectangles
    """

    logger.info(f"calculate_overlap for {face_bb_group}")

    try:
        face_one, face_two = face_bb_group

        intersection = face_two.intersection(face_one)

        return (
            intersection.area / sum([face_one.area, face_two.area]),
            intersection,
        )
    except ValueError:
        return (0, None)


def transform_rectangle(
    face_center: tuple[int, int],
    output: tuple[int, int],
    frame: tuple[int, int],
    face_scale: float,
    aspect_ratio=9 / 16,
) -> tuple[list[int], Polygon]:
    """
    Calculate face bounding box mapped to output respecting frame dimension

    Returns

    `[(x1, y1), (x2, y1), (x2, y2), (x1, y2)]`
    """
    output_width, output_height = output
    frame_width, frame_height = frame
    scale = max([10 / face_scale, frame_width, output_width]) / frame_width

    # Adjust scale factor for rectangle size
    half_width = output_width * scale * 1.3
    half_height = output_height / 2

    cx, cy = face_center

    corners = [
        (cx - half_width, cy - half_height),  # Bottom-left
        (cx + half_width, cy - half_height),  # Bottom-right
        (cx + half_width, cy + half_height),  # Top-right
        (cx - half_width, cy + half_height),  # Top-left
    ]

    rectangle = Polygon(corners)

    return [int(i) for i in rectangle.bounds], rectangle


def draw_circle(frame, coordinates, color):
    cv2.circle(
        frame,
        coordinates,
        radius=5,
        color=color,
        thickness=-1,
    )


def draw_rectangle(frame, coordinateTL, coordinateBR, color):
    cv2.rectangle(
        frame,
        coordinateTL,
        coordinateBR,
        color,
        cv2.FILLED,
    )


def draw_text(
    frame,
    text,
    coordinates,
    color,
):
    cv2.putText(
        frame,
        text,
        coordinates,
        cv2.QT_FONT_NORMAL,
        0.7,
        color,
        1,
    )


def normalized_face(frame_list):
    def average(face_list):
        width_list = []

        for face in face_list:
            if "x1" in face:
                width_list.append(face["x2"] - face["x1"])

        return max(width_list + [0])

    face_width_list = sum([average(frame["faces"]) for frame in frame_list] + [0])
    average_face = face_width_list / len(frame_list)

    def map_face(face_list, average):
        face_with_scale = []

        for face in face_list:
            if "x1" in face:
                face["scale"] = 1 - math.sin((face["x2"] - face["x1"]) / average)
                face_with_scale.append(face)

        return face_with_scale

    return [
        {
            "frame": frame_info["frame"],
            "faces": map_face(frame_info["faces"], average_face),
        }
        for frame_info in frame_list
    ]


media_list = [
    (
        "/home/sky/Downloads/76100122-cc66-472e-b1aa-5ae3355e27d6-2024-06-26-09-07-06.json",
        "video_data/76100122-cc66-472e-b1aa-5ae3355e27d6-2024-06-26-09-07-06.mp4",
    ),
    (
        "/home/sky/Downloads/c1b40654-bd35-42a3-837b-41afad797747-2024-06-27-05-19-04.json",
        "video_data/c1b40654-bd35-42a3-837b-41afad797747-2024-06-27-05-19-04.mp4",
    ),
    (
        "/home/sky/Downloads/c613877f-d6ed-4115-8835-fd2c5cc27edd-2024-06-27-05-19-04.json",
        "video_data/c613877f-d6ed-4115-8835-fd2c5cc27edd-2024-06-27-05-19-04.mp4",
    ),
    (
        "/home/sky/Downloads/c5a104be-5624-41a6-bb91-8bc0437086fd-2024-06-26-09-09-23.json",
        "video_data/c5a104be-5624-41a6-bb91-8bc0437086fd-2024-06-26-09-09-23.mp4",
    ),
    (
        "/home/sky/Downloads/74e912f2-f77e-484c-8ce1-526ccbe7f3b0-2024-06-26-07-07-15.json",
        "/home/sky/Downloads/74e912f2-f77e-484c-8ce1-526ccbe7f3b0-2024-06-26-07-07-15.mp4",
    ),
]

media_index = 4

video_json_file = open(media_list[media_index][0])
frame_info_list = normalized_face(json.load(video_json_file))

video_raw = cv2.VideoCapture(media_list[media_index][1])

OUTPUT_RESOLUTION = (1080, 1920)


video_frames = int(video_raw.get(cv2.CAP_PROP_FRAME_COUNT))
error_string = f"Video frame count ({video_frames}) should match with json frames ({len(frame_info_list)})"
assert video_frames == len(frame_info_list), error_string

start_frame = 1000

video_raw.set(cv2.CAP_PROP_POS_FRAMES, start_frame)

frame_index = start_frame
while video_raw.isOpened():
    ret, frame = video_raw.read()

    frame_cpy = frame.copy()
    frame_overlay = frame.copy()

    frame_info = frame_info_list[frame_index]
    frame_width = int(video_raw.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(video_raw.get(cv2.CAP_PROP_FRAME_HEIGHT))

    logger.info(f"processing frame {frame_index}:{frame_info}")

    width, height = (OUTPUT_RESOLUTION[0] // 2, OUTPUT_RESOLUTION[1] // 2)
    aspect_ratio = 9 / 16

    face_list = sorted(
        [f for f in frame_info["faces"] if "x1" in f], key=lambda x: x.get("x")
    )

    face_bb_list: list[Polygon] = []

    for face_index, face in enumerate(face_list):
        logger.debug(f"adding face {face}")

        face_center = (int(face.get("x")), int(face.get("y")))
        face_rect = [
            (int(face.get("x1")), int(face.get("y1"))),
            (int(face.get("x2")), int(face.get("y2"))),
        ]

        if DEBUGGING:
            cv2.circle(frame, face_center, 5, (0, 0, 0), -1)

        if DEBUGGING:
            logger.info(f"frame dimension {frame_width}x{frame_height}")
            logger.info(f"output dimension {width}x{height}")

        face_output, face_polygon = transform_rectangle(
            face_center=face_center,
            output=(width, height),
            frame=(frame_width, frame_height),
            face_scale=face["scale"],
        )

        if DEBUGGING:
            logger.debug(f"{face_output}")
            draw_rectangle(
                frame,
                (face_output[0], face_output[1]),
                (face_output[2], face_output[3]),
                COLORS[face_index],
            )

        draw_text(
            frame,
            f"Face {face_index} Scale: {face['scale']}",
            (20, 60 + face_index % len(face_list) * 20),
            (255, 255, 255),
        )
        face_bb_list.append(face_polygon)

    face_bb_chunk = [face_bb_list[i : i + 2] for i in range(0, len(face_bb_list), 2)]

    face_bb_overlap = [
        calculate_overlap(face_bb_group, (width, height))
        for face_bb_group in face_bb_chunk
    ]

    if DEBUGGING:
        alpha = 0.6
        draw_text(
            frame, f"frame: {frame_width}x{frame_height}", (20, 40), (255, 255, 255)
        )

        for overlap, polygon in face_bb_overlap:
            logger.info(f"{overlap}")
            if overlap <= 0 or polygon is None:
                logger.warning("no overlapping found")
                continue

            coordinates = tuple(polygon.exterior.coords)

            draw_rectangle(
                frame,
                (int(coordinates[0][0]), int(coordinates[0][1])),
                (int(coordinates[2][0]), int(coordinates[2][1])),
                COLORS[5],
            )
            draw_text(frame, f"Overlap: {overlap:.2f}", (20, 20), (255, 255, 255))

        frame_overlay = cv2.addWeighted(frame, 1 - alpha, frame_cpy, alpha, gamma=0)
        cv2.imshow("main", frame_overlay)

        if cv2.waitKey(0) == 27:
            exit(0)

    frame_index += 1

cv2.destroyAllWindows()