mozaikukun.py

import math
import os
from glob import glob
import threading
from multiprocessing import Process, Queue
import json
import time

from PIL.Image import Image
from ultralytics import YOLO
from PIL import Image, ImageDraw
from typing import Dict, Tuple

BLOCK_SIZE_RATIO = 100
MARGIN_FACTOR = 3
MARGIN_EXTRA = 20
DEVICE='cpu' # Specify the device number when using CUDA. Example: DEVICE='0'


def calculate_pixel_block_and_margin(image: Image.Image) -> Tuple[int, int]:
    """
    Calculate the pixel block size and margin based on image dimension.
    """
    block_size = math.ceil(max(image.width, image.height) / BLOCK_SIZE_RATIO)
    margin = block_size * MARGIN_FACTOR + MARGIN_EXTRA
    return block_size, margin


def resize_bounding_box(bounding_box: Dict[str, int], margin: int, block_size: int) -> Dict[str, int]:
    """
    Resize the bounding box by adding the margin.
    """
    bounding_box["x1"] = int((bounding_box["x1"] - margin) / block_size) * block_size
    bounding_box["x2"] = int(math.ceil((bounding_box["x2"] + margin) / block_size) * block_size)

    bounding_box["y1"] = int((bounding_box["y1"] - margin) / block_size) * block_size
    bounding_box["y2"] = int(math.ceil((bounding_box["y2"] + margin) / block_size) * block_size)

    return bounding_box


def create_transparent_mask(segment_box: Dict[str, int], segments: Dict[str, list]) -> tuple[Image, Image]:
    """
    Create a transparent mask for image segmentation.
    """
    empty_mask = Image.new("RGBA", (
        int(segment_box["x2"] - segment_box["x1"]), int(segment_box["y2"] - segment_box["y1"])), (0, 0, 0, 0))
    mask = Image.new('RGBA', (
        int(segment_box["x2"] - segment_box["x1"]), int(segment_box["y2"] - segment_box["y1"])), (0, 0, 0, 0))
    mask_draw = ImageDraw.Draw(mask)

    adjusted_x_segments = [math.ceil(segment_x) - segment_box["x1"] for segment_x in segments['x']]
    adjusted_y_segments = [math.ceil(segment_y) - segment_box["y1"] for segment_y in segments['y']]

    polygon_points = list(zip(adjusted_x_segments, adjusted_y_segments))
    mask_draw.polygon(polygon_points, fill=(255, 255, 255, 255))

    for i in range(len(polygon_points) - 1):
        mask_draw.line([polygon_points[i], polygon_points[i + 1]], fill=(255, 255, 255, 255), width=5)

    return mask, empty_mask


def apply_mask_on_image(cropped_region: Image.Image, segment_box: Dict[str, int], mask: Image.Image,
                        empty_mask: Image.Image) -> Image.Image:
    """
    Apply the created mask on the image.
    """
    segment_region = cropped_region.crop(
        (segment_box["x1"], segment_box["y1"], segment_box["x2"], segment_box["y2"]))
    empty_mask.paste(segment_region)

    masked_image = Image.new("RGBA", empty_mask.size)
    masked_image.paste(empty_mask, mask=mask)
    return masked_image


def generate_mosaic(masked_image: Image.Image, block_size: int) -> Image.Image:
    """
    Generate a mosaic image from the masked image and block size.
    """
    small_masked_image = masked_image.resize(
        (int(masked_image.size[0] // block_size), int(masked_image.size[1] // block_size)), Image.BILINEAR)
    mosaic_masked_image = small_masked_image.resize(masked_image.size, Image.NEAREST)
    return mosaic_masked_image


def image_detection_worker(process_id: int, input_queue: Queue, result_queue: Queue):
    print(f"start subprocess {process_id}")
    object_detector = YOLO("yolov8x.pt")
    segmenter = YOLO("myseg9.pt")
    while True:
        img, name = input_queue.get()
        print(f"{process_id}: get {name}")
        start_time = time.time()
        result_queue.put((process_and_analyze_image(img, object_detector, segmenter), name))
        end_time = time.time()
        elapsed_time = end_time - start_time
        print(f"{process_id}: put ({elapsed_time} sec) {name}")


def process_and_analyze_image(image: Image.Image, object_detector: YOLO, segmenter: YOLO) -> Dict:
    """
    Process a single image, analyze it and save the result.
    """
    original_image = image.convert("RGBA")

    block_size, margin = calculate_pixel_block_and_margin(original_image)

    detection_results = object_detector(original_image, save=False, device=DEVICE, project="yolov8x", name="pname1",
                                        verbose=False)

    result = {}
    sensitive_areas = ["pussy", "penis", "sex"]

    for detection in detection_results:
        for detected_object in json.loads(detection.tojson()):
            if detected_object["name"] != "person":
                continue
            bounding_box = resize_bounding_box(detected_object["box"], margin, block_size)

            cropped_region = original_image.crop(
                (bounding_box["x1"], bounding_box["y1"], bounding_box["x2"], bounding_box["y2"]))

            segmentation_results = segmenter(cropped_region, save=False, device=DEVICE, project="myseg2", name="pname2",
                                             verbose=False)
            for segmentation in segmentation_results:
                for segmented_object in json.loads(segmentation.tojson()):
                    name = segmented_object["name"]
                    #if name not in sensitive_areas:
                    #    continue
                    segments = segmented_object["segments"]
                    segment_box = resize_bounding_box(segmented_object["box"], margin, block_size)

                    if len(segments.get("x", [])) <= 2:
                        continue
                    if len(segments.get("y", [])) <= 2:
                        continue

                    mask, empty_mask = create_transparent_mask(segment_box, segments)
                    masked_image = apply_mask_on_image(cropped_region, segment_box, mask, empty_mask)
                    mosaic_masked_image = generate_mosaic(masked_image, block_size)

                    final_image = Image.new('RGBA', original_image.size)
                    mosaic_position = (
                        int(bounding_box["x1"] + segment_box["x1"]), int(bounding_box["y1"] + segment_box["y1"]))
                    final_image.paste(mosaic_masked_image, mosaic_position)
                    if name not in result:
                        result[name] = []
                    result[name].append(increase_image_opacity(final_image))

    return result


def increase_image_opacity(img: Image.Image, rate: int = 10) -> Image.Image:
    img = img.convert("RGBA")  # ensure image has alpha channel

    datas = img.getdata()
    new_data = []
    for item in datas:
        if item[3] != 0:
            new_data.append((item[0], item[1], item[2], min(255, int(item[3] * rate))))
        else:
            new_data.append(item)  # leave fully transparent pixel as it is
    img.putdata(new_data)
    return img


def main():
    print("Start processing...")
    worker_count = 3

    worker_processes = []
    worker_queue = {}
    task_queue = Queue()
    result_queue = Queue()

    for worker_number in range(worker_count):
        worker_queue[worker_number] = Queue()
        p = Process(
            target=image_detection_worker,
            args=(worker_number, task_queue, result_queue, worker_queue[worker_number]))
        p.start()
        worker_processes.append(p)

    output_dir = "output"

    os.makedirs(output_dir, exist_ok=True)

    for image_path in glob("input/*.jpg"):
        task_queue.put((Image.open(image_path), image_path, {}))

    sensitive_areas = ["pussy", "penis", "sex"]

    while True:
        result_data, image_name = result_queue.get()
        for sensitive_area in sensitive_areas:
            image_number = 0
            for result_image in result_data[sensitive_area]:
                new_image_filename = (f"{os.path.splitext(os.path.basename(image_name))[0]}_"
                                      f"{sensitive_area}_{image_number}.png")
                image_save_path = os.path.join(output_dir, new_image_filename)
                result_image.save(image_save_path)
                image_number += 1


if __name__ == '__main__':
    main()