Add application file

app.py

@@ -0,0 +1,49 @@
+import numpy as np
+import gradio as gr
+from PIL import Image
+
+import mozaikukun as moza
+from ultralytics import YOLO
+
+object_detector = YOLO("yolov8x.pt")
+segmenter = YOLO("myseg3.pt")
+
+
+def mosaic_process(input_img, pussy, penis, sex, anus, nipple):
+    img = Image.fromarray(np.uint8(input_img))
+    img = img.convert("RGBA")
+
+    process_mode = {
+        "pussy": pussy,
+        "penis": penis,
+        "sex": sex,
+        "anus": anus,
+        "nipple": nipple,
+    }
+
+    result = moza.process_and_analyze_image(img, object_detector, segmenter)
+    for key in result.keys():
+        if key not in process_mode:
+            continue
+
+        if process_mode.get(key) == "mosaic":
+            for mosaic_img in result[key]:
+                mosaic_img = mosaic_img.convert("RGBA")
+                img = Image.alpha_composite(img, mosaic_img)
+
+    return img
+
+
+demo = gr.Interface(
+    fn=mosaic_process,
+    inputs=[
+        gr.Image(),
+        gr.Radio(choices=['raw', 'mosaic'], value='mosaic'),
+        gr.Radio(choices=['raw', 'mosaic'], value='mosaic'),
+        gr.Radio(choices=['raw', 'mosaic'], value='mosaic'),
+        gr.Radio(choices=['raw', 'mosaic'], value='raw'),
+        gr.Radio(choices=['raw', 'mosaic'], value='raw'),
+    ],
+    outputs=["image"])
+
+demo.launch(server_name='0.0.0.0')

launch.py

@@ -0,0 +1,139 @@
+import importlib.util
+import os
+import shlex
+import subprocess
+import sys
+
+commandline_args = os.environ.get("COMMANDLINE_ARGS", "")
+sys.argv += shlex.split(commandline_args)
+
+python = sys.executable
+git = os.environ.get("GIT", "git")
+index_url = os.environ.get("INDEX_URL", "")
+stored_commit_hash = None
+skip_install = False
+
+
+def run(command, desc=None, errdesc=None, custom_env=None):
+    if desc is not None:
+        print(desc)
+
+    result = subprocess.run(
+        command,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        shell=True,
+        env=os.environ if custom_env is None else custom_env,
+    )
+
+    if result.returncode != 0:
+        message = f"""{errdesc or 'Error running command'}.
+Command: {command}
+Error code: {result.returncode}
+stdout: {result.stdout.decode(encoding="utf8", errors="ignore") if len(result.stdout)>0 else '<empty>'}
+stderr: {result.stderr.decode(encoding="utf8", errors="ignore") if len(result.stderr)>0 else '<empty>'}
+"""
+        raise RuntimeError(message)
+
+    return result.stdout.decode(encoding="utf8", errors="ignore")
+
+
+def check_run(command):
+    result = subprocess.run(
+        command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True
+    )
+    return result.returncode == 0
+
+
+def is_installed(package):
+    try:
+        spec = importlib.util.find_spec(package)
+    except ModuleNotFoundError:
+        return False
+
+    return spec is not None
+
+
+def commit_hash():
+    global stored_commit_hash
+
+    if stored_commit_hash is not None:
+        return stored_commit_hash
+
+    try:
+        stored_commit_hash = run(f"{git} rev-parse HEAD").strip()
+    except Exception:
+        stored_commit_hash = "<none>"
+
+    return stored_commit_hash
+
+
+def run_pip(args, desc=None):
+    if skip_install:
+        return
+
+    index_url_line = f" --index-url {index_url}" if index_url != "" else ""
+    return run(
+        f'"{python}" -m pip {args} --prefer-binary{index_url_line}',
+        desc=f"Installing {desc}",
+        errdesc=f"Couldn't install {desc}",
+    )
+
+
+def run_python(code, desc=None, errdesc=None):
+    return run(f'"{python}" -c "{code}"', desc, errdesc)
+
+
+def extract_arg(args, name):
+    return [x for x in args if x != name], name in args
+
+
+def prepare_environment():
+    commit = commit_hash()
+
+    print(f"Python {sys.version}")
+    print(f"Commit hash: {commit}")
+
+    torch_command = os.environ.get(
+        "TORCH_COMMAND",
+        "pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118",
+    )
+
+    sys.argv, skip_install = extract_arg(sys.argv, "--skip-install")
+    if skip_install:
+        return
+
+    sys.argv, reinstall_torch = extract_arg(sys.argv, "--reinstall-torch")
+    ngrok = "--ngrok" in sys.argv
+
+    if reinstall_torch or not is_installed("torch") or not is_installed("torchaudio"):
+        run(
+            f'"{python}" -m {torch_command}',
+            "Installing torch and torchaudio",
+            "Couldn't install torch",
+        )
+
+    if not is_installed("pyngrok") and ngrok:
+        run_pip("install pyngrok", "ngrok")
+
+    run(
+        f'"{python}" -m pip install -r requirements.txt',
+        desc=f"Installing requirements",
+        errdesc=f"Couldn't install requirements",
+    )
+
+
+def start():
+    os.environ["PATH"] = (
+        os.path.join(os.path.dirname(__file__), "bin")
+        + os.pathsep
+        + os.environ.get("PATH", "")
+    )
+    subprocess.run(
+        [python, "webui.py", *sys.argv[1:]],
+    )
+
+
+if __name__ == "__main__":
+    prepare_environment()
+    start()

mozaikukun.py

@@ -0,0 +1,208 @@
+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("myseg3.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()

package.txt

@@ -0,0 +1,3 @@
+libgl1-mesa-glx
+libgl1-mesa-dri
+