insert.py

# CS5330 Lab 4
# This file includes functions for inserting a person into a pair of stereoscopic images

import os
from PIL import Image, ImageEnhance, ImageFilter
from enum import Enum
from segmentation import extract_person

# Define an enum class for the background (background must be chosen from the three provided images)
class Background(Enum):
    NATURE = os.path.join('background', 'spatial_sbs_2024-10-10_21-57-07-266Z.jpg')
    NEU = os.path.join('background', 'sbs_neu.jpg')
    GASTOWN = os.path.join('background', 'side_by_side_steam_clock.jpg')

    def get_background(self):
        return Image.open(self.value)


# Define an enum class for disparity (disparity must be chosen from close, medium or far)
class Disparity(Enum):
    CLOSE = 30
    MED = 50
    FAR = 100


# Add person onto background image
def overlay_person(background, person, x_pos, y_pos):
    background = background.convert("RGBA")
    person = person.convert("RGBA")

    # Match lighting and color between the background and the person
    match_lighting_and_color(background, person)

    # Paste the person onto the background
    background.paste(person, (x_pos, y_pos), person)

    return background


# Scale the person according to disparity to enhance depth illusion
def scale_person(person, disparity):
    # Define a scaling factor based on disparity
    if disparity == Disparity.FAR:
        scaling_factor = 0.75
    elif disparity == Disparity.MED:
        scaling_factor = 0.9
    else:
        scaling_factor = 1

    new_width = int(person.size[0] * scaling_factor)
    new_height = int(person.size[1] * scaling_factor)
    scaled_person = person.resize((new_width, new_height))
    return scaled_person


# Function to match color and lighting
def match_lighting_and_color(background, person):
    # Match brightness
    brightness_factor = get_brightness_factor(background, person)
    person = ImageEnhance.Brightness(person).enhance(brightness_factor)
    
    # Match contrast
    contrast_factor = get_contrast_factor(background, person)
    person = ImageEnhance.Contrast(person).enhance(contrast_factor)

    return person


# Calculate the ratio between background brightness and person brightness
def get_brightness_factor(background, person):
    # Get the average brightness of both images
    background_avg_brightness = calculate_avg_brightness(background)
    person_avg_brightness = calculate_avg_brightness(person)
    
    # Return the ratio of the brightness levels
    return background_avg_brightness / person_avg_brightness


# Calculate the ratio between background contrast and person contrast
def get_contrast_factor(background, person):
    # Get the average contrast of both images
    background_avg_contrast = calculate_avg_contrast(background)
    person_avg_contrast = calculate_avg_contrast(person)
    
    # Return the ratio of the contrast levels
    return background_avg_contrast / person_avg_contrast

# Calculate the average brightness of an image
def calculate_avg_brightness(image):
    grayscale_image = image.convert('L')
    pixels = list(grayscale_image.getdata())
    return sum(pixels) / len(pixels)

# Calculate the average contrast of an image
def calculate_avg_contrast(image):
    grayscale_image = image.convert('L')
    contrast_image = grayscale_image.filter(ImageFilter.FIND_EDGES)
    pixels = list(contrast_image.getdata())
    return sum(pixels) / len(pixels)


# Place the person on the pair of stereoscopic images
# Depth perception can be manipulated by adjusting disparity
# background: an enum, must be chosen from the three provided background
# person: the extracted person as a transparent png
# disparity: an enum, must be chosen from the close, medium or far
# The bottom of the person is always going to align with the bottom of the background
def insert_person_in_pair(background: Background, person, disparity: Disparity):
    # Get the background image according to the choice
    background_image = background.get_background()
    # Split the stereoscopic image into left and right images
    width, height = background_image.size
    left_image = background_image.crop((0, 0, width // 2, height))
    right_image = background_image.crop((width // 2, 0, width, height))
    person_height = person.size[1]

    # Define an appropriate position for the person according to background and depth perception
    # Ensure that the person appears at an appropriate position and blends in the background naturally
    y_pos = height - person_height
    if background == Background.NATURE:
        x_pos = 1750
        if disparity == Disparity.MED:
            x_pos += 200
            y_pos -= 25
        elif disparity == Disparity.FAR:
            x_pos += 300
            y_pos -= 200
    elif background == Background.NEU:
        x_pos = 1500
        if disparity == Disparity.MED:
            x_pos += 200
            y_pos += 50
        elif disparity == Disparity.FAR:
            x_pos += 300
            y_pos += 100
    else:
        x_pos = 350
        if disparity == Disparity.MED:
            x_pos += 300
            y_pos -= 25
        elif disparity == Disparity.FAR:
            x_pos += 400
            y_pos -= 50
    
    # Scale the person to enhance depth illusion
    person = scale_person(person, disparity)

    # Calculate position for the person in the right image based on disparity
    x_pos_right = x_pos - disparity.value

    # Overlay the person onto left and right images at calculated positions
    left_with_person = overlay_person(left_image.copy(), person, x_pos, y_pos)
    right_with_person = overlay_person(right_image.copy(), person, x_pos_right, y_pos)

    return left_with_person, right_with_person


# ==========Test case==========
# person_image = extract_person(Image.open('test_img/mona_lisa.jpg'))

# left, right = insert_person_in_pair(
#     background=Background.NATURE,
#     person=person_image,
#     disparity=Disparity.MED
# )