add app files

app.py

@@ -0,0 +1,22 @@
+import uvicorn
+import threading
+import gradio as gr
+from app.main import app as fastapi_app
+
+# run FastAPI in a background thread
+def run_fastapi():
+    uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)
+
+thread = threading.Thread(target=run_fastapi, daemon=True)
+thread.start()
+
+# Gradio needs a UI to stay alive — minimal dummy interface
+demo = gr.Interface(
+    fn=lambda x: "Emosaic backend is running!",
+    inputs=gr.Textbox(label="ping"),
+    outputs=gr.Textbox(label="status"),
+    title="🌸 Emosaic Backend",
+    description="FastAPI backend for Emosaic. Use the /generate endpoint."
+)
+
+demo.launch(server_port=7860)

app/color_detector.py

@@ -0,0 +1,73 @@
+import numpy as np
+
+
+COLOR_CENTROIDS_LAB = {
+    "dark":    [10,   0,    0  ],
+    "light":   [95,   0,    0  ],
+    "red":     [40,   55,   35 ],
+    "pink":    [65,   35,  -5  ],
+    "orange":  [60,   25,   45 ],
+    "brown":   [35,   15,   20 ],
+    "yellow":  [85,  -5,    60 ],
+    "green":   [50,  -40,   25 ],
+    "blue":    [45,   5,   -40 ],
+    "neutral": [70,   2,    4  ],
+}
+
+
+def rgb_to_lab(r, g, b):
+    def linearize(c):
+        return c / 12.92 if c <= 0.04045 else ((c + 0.055) / 1.055) ** 2.4
+
+    r, g, b = linearize(r), linearize(g), linearize(b)
+
+    X = r * 0.4124564 + g * 0.3575761 + b * 0.1804375
+    Y = r * 0.2126729 + g * 0.7151522 + b * 0.0721750
+    Z = r * 0.0193339 + g * 0.1191920 + b * 0.9503041
+
+    X /= 0.95047
+    Y /= 1.00000
+    Z /= 1.08883
+
+    def f(t):
+        return t ** (1/3) if t > 0.008856 else (7.787 * t + 16/116)
+
+    fx, fy, fz = f(X), f(Y), f(Z)
+
+    L = 116 * fy - 16
+    a = 500 * (fx - fy)
+    b_val = 200 * (fy - fz)
+
+    return L, a, b_val
+
+
+def get_color_group(pixel):
+    r, g, b = [x / 255.0 for x in pixel[:3]]
+    L, a, b_val = rgb_to_lab(r, g, b)
+
+    if L < 18:
+        return "dark"
+    if L > 92 and abs(a) < 5 and abs(b_val) < 5:
+        return "light"
+
+    lab = [L, a, b_val]
+    best_color = min(
+        COLOR_CENTROIDS_LAB,
+        key=lambda c: (
+            (lab[0] - COLOR_CENTROIDS_LAB[c][0]) ** 2 * 0.3 +   # weight L less
+            (lab[1] - COLOR_CENTROIDS_LAB[c][1]) ** 2 * 1.5 +   # weight a more
+            (lab[2] - COLOR_CENTROIDS_LAB[c][2]) ** 2 * 1.5     # weight b more
+        )
+    )
+
+    return best_color
+
+
+def build_color_grid(pixels):
+    color_grid = []
+    for row in pixels:
+        color_row = []
+        for pixel in row:
+            color_row.append(get_color_group(pixel))
+        color_grid.append(color_row)
+    return color_grid

app/emoji_mapper.py

@@ -0,0 +1,140 @@
+import json
+import random
+import numpy as np
+
+from app.themes import THEMES
+
+with open("data/emoji_colors.json", "r", encoding="utf-8") as f:
+    _raw_colors = json.load(f)
+
+EMOJI_COLOR_LOOKUP = {
+    char: data["lab"] for char, data in _raw_colors.items()
+}
+
+
+def rgb_to_lab_batch(pixels):
+    """
+    Vectorized RGB -> LAB conversion.
+    pixels: numpy array (N, 3) in 0-255
+    Returns: numpy array (N, 3) LAB
+    """
+    rgb = pixels.astype(np.float64) / 255.0
+
+    def linearize(c):
+        return np.where(c <= 0.04045, c / 12.92, ((c + 0.055) / 1.055) ** 2.4)
+
+    rgb_lin = linearize(rgb)
+
+    # sRGB -> XYZ matrix
+    M = np.array([
+        [0.4124564, 0.3575761, 0.1804375],
+        [0.2126729, 0.7151522, 0.0721750],
+        [0.0193339, 0.1191920, 0.9503041],
+    ])
+
+    xyz = rgb_lin @ M.T
+
+    xyz[:, 0] /= 0.95047
+    xyz[:, 1] /= 1.00000
+    xyz[:, 2] /= 1.08883
+
+    def f(t):
+        return np.where(t > 0.008856, np.cbrt(t), 7.787 * t + 16 / 116)
+
+    fxyz = f(xyz)
+
+    L = 116 * fxyz[:, 1] - 16
+    a = 500 * (fxyz[:, 0] - fxyz[:, 1])
+    b = 200 * (fxyz[:, 1] - fxyz[:, 2])
+
+    return np.stack([L, a, b], axis=1)
+
+
+def build_candidate_lab_matrix(candidates):
+    """
+    candidates: list of emoji chars
+    Returns: (valid_chars, lab_matrix) where lab_matrix is (M, 3)
+    """
+    valid_chars = []
+    labs = []
+    for char in candidates:
+        lab = EMOJI_COLOR_LOOKUP.get(char)
+        if lab is not None:
+            valid_chars.append(char)
+            labs.append(lab)
+
+    if not valid_chars:
+        # fallback — no color data, just return first candidate with dummy lab
+        return [candidates[0]], np.array([[50.0, 0.0, 0.0]])
+
+    return valid_chars, np.array(labs)
+
+
+# cache candidate LAB matrices per (theme, category) so we don't rebuild every request
+_candidate_cache = {}
+
+
+def get_candidate_matrix(theme_name, label, cell_type):
+    key = (theme_name, label, cell_type)
+    if key not in _candidate_cache:
+        theme = THEMES.get(theme_name, THEMES["pastel"])
+        if cell_type == "segment":
+            candidates = theme.get(label, theme.get("background", ["⬜"]))
+        else:
+            candidates = theme.get(label, theme["neutral"])
+        _candidate_cache[key] = build_candidate_lab_matrix(candidates)
+    return _candidate_cache[key]
+
+
+def build_emoji_grid(segment_grid, theme_name, chaos):
+    """
+    Vectorized version: groups cells by (cell_type, label) so we batch-process
+    all pixels needing the same candidate pool at once.
+    """
+    theme_name = theme_name.lower()
+    chaos = max(0.0, min(1.0, chaos))
+
+    H = len(segment_grid)
+    W = len(segment_grid[0]) if H else 0
+
+    # flatten grid into arrays for batch processing
+    flat_cells = [cell for row in segment_grid for cell in row]
+    cell_types = [c[0] for c in flat_cells]
+    labels = [c[1] for c in flat_cells]
+    pixels = np.array([c[2] for c in flat_cells])  # (N, 3)
+
+    N = len(flat_cells)
+    result = [None] * N
+
+    # group indices by (cell_type, label) so each group shares a candidate pool
+    groups = {}
+    for i in range(N):
+        key = (cell_types[i], labels[i])
+        groups.setdefault(key, []).append(i)
+
+    for (cell_type, label), indices in groups.items():
+        valid_chars, lab_matrix = get_candidate_matrix(theme_name, label, cell_type)
+
+        idx_array = np.array(indices)
+        group_pixels = pixels[idx_array]  # (G, 3)
+
+        target_labs = rgb_to_lab_batch(group_pixels)  # (G, 3)
+
+        # pairwise squared distance: (G, 1, 3) - (1, M, 3) -> (G, M, 3) -> sum -> (G, M)
+        diffs = target_labs[:, None, :] - lab_matrix[None, :, :]
+        dists = np.sum(diffs ** 2, axis=2)
+
+        best_idx = np.argmin(dists, axis=1)  # (G,)
+
+        for j, i in enumerate(indices):
+            if chaos > 0 and random.random() < chaos:
+                result[i] = random.choice(valid_chars)
+            else:
+                result[i] = valid_chars[best_idx[j]]
+
+    # reshape back into grid
+    emoji_grid = []
+    for y in range(H):
+        emoji_grid.append(result[y * W:(y + 1) * W])
+
+    return emoji_grid

app/image_processor.py

@@ -0,0 +1,39 @@
+import numpy as np
+from PIL import Image
+
+
+def load_image(file):
+    image = Image.open(file)
+
+    # Convert everything to RGB
+    image = image.convert("RGB")
+
+    return image
+
+
+def resize_image(image, resolution):
+    width, height = image.size
+
+    # Clamp resolution
+    resolution = max(40, min(300, resolution))  
+
+    aspect_ratio = width / height
+
+    if width >= height:
+        new_width = resolution
+        new_height = int(resolution / aspect_ratio)
+    else:
+        new_height = resolution
+        new_width = int(resolution * aspect_ratio)
+
+    resized = image.resize(
+        (new_width, new_height),
+        Image.LANCZOS
+    )
+
+    return resized
+
+
+def image_to_array(image):
+
+    return np.array(image)

app/main.py

@@ -0,0 +1,17 @@
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from app.routes import router
+
+app = FastAPI(
+    title="Emoji Mosaic API",
+    version="1.0"
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["http://localhost:5173"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+app.include_router(router)

app/routes.py

@@ -0,0 +1,134 @@
+import time
+
+from fastapi import (
+    APIRouter,
+    UploadFile,
+    File,
+    Form,
+    HTTPException
+)
+
+from app.image_processor import (
+    load_image,
+    resize_image,
+    image_to_array
+)
+
+from app.emoji_mapper import (
+    build_emoji_grid
+)
+
+from app.themes import THEMES
+from ml.segmentor import segment_image
+
+router = APIRouter()
+
+
+@router.post("/generate")
+async def generate(
+
+    image: UploadFile = File(...),
+
+    theme: str = Form(...),
+
+    resolution: int = Form(...),
+
+    chaos: float = Form(...)
+
+):
+
+    start = time.time()
+
+    # -----------------------------
+    # Validate uploaded file
+    # -----------------------------
+
+    if not image.content_type.startswith("image/"):
+
+        raise HTTPException(
+            status_code=400,
+            detail="Only image files are allowed."
+        )
+
+    # -----------------------------
+    # Validate theme
+    # -----------------------------
+
+    theme = theme.lower()
+
+    if theme not in THEMES:
+
+        raise HTTPException(
+            status_code=400,
+            detail=f"Theme '{theme}' not found."
+        )
+
+    # -----------------------------
+    # Validate resolution
+    # -----------------------------
+
+    if resolution < 20 or resolution > 100:
+
+        raise HTTPException(
+            status_code=400,
+            detail="Resolution must be between 20 and 100."
+        )
+
+    # -----------------------------
+    # Validate chaos
+    # -----------------------------
+
+    if chaos < 0 or chaos > 1:
+
+        raise HTTPException(
+            status_code=400,
+            detail="Chaos must be between 0 and 1."
+        )
+
+    # -----------------------------
+    # Process image
+    # -----------------------------
+
+    image_obj = load_image(image.file)
+    internal_resolution = int(40 + (resolution - 20) * (300 - 40) / (100 - 20))
+    resized = resize_image(image_obj, internal_resolution)
+
+    pixels = image_to_array(resized)
+
+    # -----------------------------
+    # Segment image (CNN + fallback)
+    # -----------------------------
+
+    segment_grid = segment_image(resized, pixels)
+
+    # -----------------------------
+    # Build emoji grid
+    # -----------------------------
+
+    emoji_grid = build_emoji_grid(
+        segment_grid,
+        theme,
+        chaos
+    )
+
+    end = time.time()
+
+    print(f"Processing Time: {end-start:.4f} seconds")
+
+    return {
+
+        "success": True,
+
+        "theme": theme,
+
+        "resolution": resolution,
+
+        "chaos": chaos,
+
+        "rows": len(emoji_grid),
+
+        "cols": len(emoji_grid[0]),
+
+        "grid": emoji_grid
+
+    }

app/themes.py

@@ -0,0 +1,174 @@
+PASTEL = {
+
+    # ── segment keys ──
+    "hair": [
+        "🌸", "🎀", "💗", "✨", "🩷", "🍡", "🧁", "🍥", "🌷", "🦄",
+        "🍬", "🪞", "🪄", "🎐", "🌺", "🍓", "🩰", "💞", "🐚",
+    ],
+    "skin": [
+        "🍑", "💛", "🌼", "🤍", "🫐", "🍈", "🥛", "🍮", "🧈", "🌻",
+        "🐤", "🍯", "🌾", "🥮", "🍠", "🍞", "🥐", "🧇",
+    ],
+    "clothes": [
+        "🎀", "🩷", "🧁", "💜", "👛", "👗", "👒", "🩱", "🧶", "🎒",
+        "👝", "🩲", "💄", "💍", "🛍️", "🪡", "🧵",
+    ],
+    "water": [
+        "🫧", "🩵", "🌊", "💙", "💧", "🧊", "🚿", "🪣", "🐬", "🦢",
+        "🛁", "🥤", "🧴",
+    ],
+    "plants": [
+        "🌿", "🍀", "🌱", "🌸", "🌷", "🍃", "🌹", "🪴", "🌵", "🌳",
+        "🥀", "🌾", "🍄", "🌼",
+    ],
+    "ground": [
+        "🤎", "🍪", "🧸", "🪵", "🍂", "🥨", "🌰", "🪨", "🍫", "🦴",
+        "🥖", "🧺",
+    ],
+    "background": [
+        "🤍", "☁️", "💭", "✨", "🌙", "⭐", "🪽", "🤿", "🎐", "🫐",
+        "🪞", "🦢",
+    ],
+
+    # ── color fallback keys ──
+    "pink": ["🌸", "🎀", "🩷", "🍡", "🩰", "🌷", "💐", "🦩", "🧁", "🍓"],
+    "red": ["❤️", "🍓", "🌹", "🍒", "🍎", "🟥", "🍅", "🧧"],
+    "orange": ["🍊", "🧡", "🦊", "🌅", "🍑", "🥕", "🦁", "🟧"],
+    "brown": ["🧸", "🤎", "🍪", "🪵", "🦴", "🥨", "🌰", "🐻", "🍫"],
+    "yellow": ["🌼", "⭐", "🐥", "🌻", "🍋", "🧀", "🟨", "🍌", "🌕"],
+    "green": ["🌿", "🍀", "🌱", "🥝", "🟩", "🐸", "🍏", "🌳"],
+    "blue": ["🫧", "🐳", "🩵", "🌊", "💙", "🦋", "🟦", "🐬"],
+    "light": ["☁️", "🤍", "✨", "🩶", "🫧", "🪽", "🦢"],
+    "dark": ["🌑", "🖤", "🌌", "🦇", "🐈‍⬛", "⚫"],
+    "neutral": ["🤍", "🩶", "💭", "🪨", "🐚", "🍄"],
+}
+
+
+OCEAN = {
+
+    # ── segment keys ──
+    "hair": [
+        "🪸", "🐚", "🌊", "🦀", "🐙", "🌀", "🫧", "🧜‍♀️", "🦑", "🪼",
+        "🧵", "🌫️",
+    ],
+    "skin": [
+        "🐠", "🌺", "⭐", "🪸", "🐡", "🦐", "🦞", "🐢", "🌴", "🍍",
+        "🦩", "🥥",
+    ],
+    "clothes": [
+        "🐬", "🌊", "🪼", "💙", "🩱", "🦈", "⚓", "🧜‍♂️", "🛟", "🪢",
+    ],
+    "water": [
+        "🌊", "🐳", "🫧", "🐬", "💧", "🌀", "🚣", "⛵", "🛶", "🐋",
+        "🦭", "🪣", "🧊", "🌧️",
+    ],
+    "plants": [
+        "🌿", "🐢", "🌴", "🌱", "🪸", "🥥", "🍃", "🌊", "🦦",
+    ],
+    "ground": [
+        "🪨", "🐚", "⚓", "🦑", "🏝️", "🦀", "🪵", "🦂", "🪱",
+    ],
+    "background": [
+        "🌌", "💎", "☁️", "🫧", "🌊", "🌙", "⭐", "🐳", "🪞", "🌫️",
+    ],
+
+    # ── color fallback keys ──
+    "pink": ["🪸", "🌺", "🩷", "🦩", "🐡", "🦐"],
+    "red": ["🦀", "🐙", "❤️", "🦞", "🌶️", "🍓"],
+    "orange": ["🐠", "🪼", "🧡", "🦞", "🥭", "🐡"],
+    "brown": ["🐚", "🪵", "🦑", "🦂", "🪱", "🦦"],
+    "yellow": ["⭐", "🐡", "🌞", "🍍", "🟡", "🐤"],
+    "green": ["🌿", "🐢", "🌴", "🥥", "🐊", "🦎"],
+    "blue": ["🌊", "🐳", "🫧", "🐬", "💙", "🐋", "🦈", "🔵"],
+    "light": ["☁️", "💎", "✨", "🫧", "🪽", "🌫️"],
+    "dark": ["🌌", "🌑", "⚓", "🦑", "🐙", "⚫"],
+    "neutral": ["🪨", "☁️", "💭", "🐚", "🦪"],
+}
+
+
+SPOOKY = {
+
+    # ── segment keys ──
+    "hair": [
+        "🕸️", "🖤", "💀", "🦇", "🧛", "🧟", "🕷️", "🐈‍⬛", "🪦", "🌑",
+    ],
+    "skin": [
+        "👻", "🕯️", "💀", "🌕", "☠️", "🧟", "🧛", "🦴", "🫥", "🪱",
+    ],
+    "clothes": [
+        "🧛", "🖤", "🎃", "⚰️", "🦇", "🪦", "👹", "🧙", "🩸", "🥀",
+    ],
+    "water": [
+        "🌊", "🔮", "🌑", "💀", "🪦", "☠️", "🧪", "🫗",
+    ],
+    "plants": [
+        "☠️", "🧪", "🕸️", "🌿", "🥀", "🍄", "🌑", "🪦",
+    ],
+    "ground": [
+        "🪦", "🦴", "🌑", "🖤", "⚰️", "🕳️", "🍂", "🪱",
+    ],
+    "background": [
+        "🌌", "👻", "🌫️", "🖤", "🌑", "🦇", "🕸️", "🪦", "⚫",
+    ],
+
+    # ── color fallback keys ──
+    "pink": ["🕸️", "🎃", "💀", "🩸", "🧠"],
+    "red": ["🩸", "🧛", "👹", "🔥", "🍷", "🟥"],
+    "orange": ["🎃", "🔥", "🦇", "🧡", "🍊"],
+    "brown": ["🪵", "🦉", "🦂", "🦴", "🤎", "🪱"],
+    "yellow": ["🕯️", "👁️", "🌕", "🟡", "🐍"],
+    "green": ["☠️", "🧪", "👻", "🧟", "🐍", "🟢"],
+    "blue": ["🌌", "🦇", "🔮", "🌑", "🔵"],
+    "light": ["👻", "✨", "🌙", "🤍", "🫥"],
+    "dark": ["🖤", "🌑", "⚰️", "⚫", "🦇", "🕳��"],
+    "neutral": ["🪦", "💀", "🌫️", "🩶", "🦴"],
+}
+
+
+Y2K = {
+
+    # ── segment keys ──
+    "hair": [
+        "💿", "🌸", "💅", "✨", "💎", "🔮", "🎀", "🪩", "💋", "🦋",
+        "🩷", "💗",
+    ],
+    "skin": [
+        "🍑", "💄", "📱", "👛", "💋", "🦋", "🍓", "🌟", "💫", "🫦",
+    ],
+    "clothes": [
+        "👙", "🩱", "👖", "🎧", "📼", "🪩", "💿", "🦋", "🔗", "👑",
+        "💍", "🛹",
+    ],
+    "water": [
+        "💧", "🫧", "🩵", "🐬", "🌊", "🪩",
+    ],
+    "plants": [
+        "🌸", "🦋", "🌺", "🌿", "✨", "🌟",
+    ],
+    "ground": [
+        "🪩", "💿", "📼", "🎮", "🕹️", "📟", "💽", "🔗",
+    ],
+    "background": [
+        "✨", "💫", "🌟", "🪩", "💜", "🩷", "🔮", "🌌", "💎",
+    ],
+
+    # ── color fallback keys ──
+    "pink": ["💗", "🩷", "💅", "🎀", "🦋", "💋", "🌸"],
+    "red": ["❤️‍🔥", "🔥", "🍓", "💋", "🟥", "🧨"],
+    "orange": ["🧡", "🔥", "🍊", "🟧", "🦊"],
+    "brown": ["🤎", "🪵", "🦴", "🟫"],
+    "yellow": ["⭐", "🌟", "🟡", "💛", "⚡"],
+    "green": ["🟢", "🍏", "💚", "👽", "🐍"],
+    "blue": ["🔵", "💙", "🪩", "🐬", "🫧", "💎"],
+    "light": ["✨", "🤍", "💫", "🩶", "💎"],
+    "dark": ["🖤", "⚫", "🌑", "🕶️"],
+    "neutral": ["🩶", "💭", "🪞", "📀"],
+}
+
+
+THEMES = {
+    "pastel": PASTEL,
+    "ocean":  OCEAN,
+    "spooky": SPOOKY,
+    "y2k":    Y2K,
+}

checkpoints/model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb2989c1a8b099c67160d5e1bfa778086bd428ce544841f2c86aac167e57206e
+size 2489050

ml/dataset.py

@@ -0,0 +1,235 @@
+import os
+import numpy as np
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from collections import Counter
+
+CLASS_NAMES = [
+    "background",   # 0 
+    "hair",         # 1
+    "skin",         # 2
+    "clothes",      # 3
+    "water",        # 4
+    "ground",       # 5
+    "plants",       # 6
+]
+
+NUM_CLASSES = 7
+
+ADE20K_TO_OURS = {
+    2:  0,   # sky → background
+    4:  5,   # floor → ground
+    9:  6,   # tree → plants
+    10: 5,   # grass → ground
+    14: 4,   # water
+    17: 0,   # table → background
+    21: 0,   # chair → background
+    26: 0,   # car → background
+    30: 6,   # bush → plants
+    64: 4,   # sea → water
+    67: 0,   # food → background
+    96: 0,   # bus → background
+}
+
+HUMAN_TO_OURS = {
+    0:  0,   # background
+    2:  1,   # hair
+    4:  3,   # upper clothes
+    6:  3,   # pants → clothes
+    7:  3,   # dress → clothes
+    11: 2,   # face → skin
+    13: 2,   # left leg → skin
+    14: 2,   # right leg → skin
+    15: 2,   # left arm → skin
+    16: 2,   # right arm → skin
+}
+
+PATCH_SIZE = 32
+STRIDE = 16
+
+
+# Dataset class — PyTorch needs this to load data in batches
+
+class PatchDataset(Dataset):
+
+    def __init__(self, samples, transform=None):
+        self.samples = samples
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        patch, label = self.samples[idx]
+
+        patch = Image.fromarray(patch)
+
+        if self.transform:
+            patch = self.transform(patch)
+
+        return patch, label
+
+
+# Patch extraction — slides window across one image+mask pair
+
+def extract_patches_from_image(image_path, mask_path, label_map):
+
+    image = np.array(Image.open(image_path).convert("RGB"))
+    mask = np.array(Image.open(mask_path))
+
+    H, W, _ = image.shape
+    samples = []
+
+    for y in range(0, H - PATCH_SIZE, STRIDE):
+        for x in range(0, W - PATCH_SIZE, STRIDE):
+
+            # crop patch from image
+            patch = image[y:y+PATCH_SIZE, x:x+PATCH_SIZE]
+
+            # crop same region from mask
+            mask_patch = mask[y:y+PATCH_SIZE, x:x+PATCH_SIZE]
+
+            # check center pixel label
+            center_y = PATCH_SIZE // 2
+            center_x = PATCH_SIZE // 2
+            raw_label = int(mask_patch[center_y, center_x])
+
+            # translate to our class system
+            our_label = label_map.get(raw_label, None)
+
+            # skip if this label isn't in our mapping
+            if our_label is None:
+                continue
+
+            samples.append((patch, our_label))
+
+    return samples
+
+
+# ADE20K loader — reads from disk
+
+def load_ade20k(root, max_images=500):
+
+    samples = []
+
+    img_dir = os.path.join(root, "ADEChallengeData2016", "images", "training")
+    mask_dir = os.path.join(root, "ADEChallengeData2016", "annotations", "training")
+
+    files = os.listdir(img_dir)
+    files = files[:max_images] #sliced dataset
+
+    print(f"ADE20K: found {len(files)} images")
+
+    for fname in files:
+
+        if not fname.endswith(".jpg"):
+            continue
+
+        img_path = os.path.join(img_dir, fname)
+
+        # mask has same name but .png
+        mask_path = os.path.join(mask_dir, fname.replace(".jpg", ".png"))
+
+        if not os.path.exists(mask_path):
+            continue
+
+        patches = extract_patches_from_image(
+            img_path,
+            mask_path,
+            ADE20K_TO_OURS
+        )
+
+        samples.extend(patches)
+
+    print(f"ADE20K: extracted {len(samples)} patches")
+    return samples
+
+
+# Human parsing loader — reads from HuggingFace
+
+def load_human_parsing(max_images=500):
+
+    samples = []
+
+    from datasets import load_dataset
+
+    print("Loading human parsing dataset from HuggingFace...")
+    ds = load_dataset("mattmdjaga/human_parsing_dataset", split="train")
+
+    print(f"Human parsing: found {max_images} images")
+
+    for item in ds.select(range(max_images)):
+
+        image = np.array(item["image"].convert("RGB"))
+        mask = np.array(item["mask"])
+
+        H, W, _ = image.shape
+
+        for y in range(0, H - PATCH_SIZE, STRIDE):
+            for x in range(0, W - PATCH_SIZE, STRIDE):
+
+                patch = image[y:y+PATCH_SIZE, x:x+PATCH_SIZE]
+                mask_patch = mask[y:y+PATCH_SIZE, x:x+PATCH_SIZE]
+
+                center_y = PATCH_SIZE // 2
+                center_x = PATCH_SIZE // 2
+                raw_label = int(mask_patch[center_y, center_x])
+
+                our_label = HUMAN_TO_OURS.get(raw_label, None)
+
+                if our_label is None:
+                    continue
+
+                samples.append((patch, our_label))
+
+    print(f"Human parsing: extracted {len(samples)} patches")
+    return samples
+
+
+# Main build function — combines both datasets
+
+def build_dataset(ade20k_root, transform=None, max_images_per_source=500, max_per_class=5000):
+
+    print("Building dataset...")
+
+    ade_samples = load_ade20k(ade20k_root, max_images=max_images_per_source)
+    human_samples = load_human_parsing(max_images=max_images_per_source)
+
+    all_samples = ade_samples + human_samples
+
+    # cap each class so no single class dominates
+    from collections import defaultdict
+    import random
+
+    bucketed = defaultdict(list)
+    for patch, label in all_samples:
+        bucketed[label].append((patch, label))
+
+    balanced = []
+    for label, items in bucketed.items():
+        random.shuffle(items)
+        balanced.extend(items[:max_per_class])
+
+    random.shuffle(balanced)
+
+    print(f"Total patches after balancing: {len(balanced)}")
+
+    counts = Counter(label for _, label in balanced)
+    for i, name in enumerate(CLASS_NAMES):
+        print(f"  {name}: {counts.get(i, 0)}")
+
+    dataset = PatchDataset(balanced, transform=transform)
+    return dataset
+
+# Transform — normalizes patches for CNN input
+
+def get_transform():
+    return transforms.Compose([
+        transforms.Resize((32, 32)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])

ml/model.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+
+
+class PatchCNN(nn.Module):
+
+    def __init__(self, num_classes=7):
+        super(PatchCNN, self).__init__()
+
+        # ── Feature Extractor ──
+        self.features = nn.Sequential(
+
+            # Block 1 — detects edges and basic textures
+            nn.Conv2d(3, 32, kernel_size=3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),          # 32x32 → 16x16
+
+            # Block 2 — detects patterns from edges
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),          # 16x16 → 8x8
+
+            # Block 3 — detects complex textures
+            nn.Conv2d(64, 128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),          # 8x8 → 4x4
+        )
+
+        # ── Classifier ──
+        self.classifier = nn.Sequential(
+            nn.Flatten(),                # (128, 4, 4) → 2048
+            nn.Linear(128 * 4 * 4, 256),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(256, num_classes)  # 256 → 7
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.classifier(x)
+        return x

ml/predict.py

@@ -0,0 +1,81 @@
+import torch
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+
+from ml.model import PatchCNN
+from ml.dataset import CLASS_NAMES, PATCH_SIZE
+
+CHECKPOINT_PATH = "checkpoints/model.pth"
+CONFIDENCE_THRESHOLD = 0.5   # below this → fall back to color
+
+# ── load model once at module level ──
+# this means it loads when FastAPI starts, not on every request
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = PatchCNN(num_classes=7)
+model.load_state_dict(torch.load(CHECKPOINT_PATH, map_location=device))
+model.eval()
+model.to(device)
+
+print(f"PatchCNN loaded on {device}")
+
+
+# ── same transform as training ──
+transform = transforms.Compose([
+    transforms.Resize((32, 32)),
+    transforms.ToTensor(),
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+])
+
+
+def predict_image(image):
+    """
+    Takes a PIL image.
+    Returns:
+        label_grid      — 2D list of class name strings
+        confidence_grid — 2D list of float confidence scores
+    """
+
+    image_np = np.array(image.convert("RGB"))
+    H, W, _ = image_np.shape
+    
+    if H < PATCH_SIZE or W < PATCH_SIZE:
+        return [], []
+
+    label_grid = []
+    confidence_grid = []
+
+    for y in range(0, H - PATCH_SIZE, PATCH_SIZE):   # stride = PATCH_SIZE (no overlap at inference)
+        label_row = []
+        conf_row = []
+
+        for x in range(0, W - PATCH_SIZE, PATCH_SIZE):
+
+            # extract patch
+            patch = image_np[y:y+PATCH_SIZE, x:x+PATCH_SIZE]
+            patch_pil = Image.fromarray(patch)
+
+            # transform → tensor
+            tensor = transform(patch_pil).unsqueeze(0).to(device)
+            # unsqueeze(0) adds batch dimension: (3,32,32) → (1,3,32,32)
+
+            # inference
+            with torch.no_grad():
+                output = model(tensor)               # (1, 7) raw logits
+                probs = torch.softmax(output, dim=1) # convert to probabilities
+                confidence, pred = probs.max(dim=1)  # get top class + its confidence
+
+            label = CLASS_NAMES[pred.item()]
+            conf = confidence.item()
+
+            label_row.append(label)
+            conf_row.append(conf)
+
+        label_grid.append(label_row)
+        confidence_grid.append(conf_row)
+
+    return label_grid, confidence_grid

ml/segmentor.py

@@ -0,0 +1,61 @@
+import numpy as np
+from PIL import Image
+from ml.predict import predict_image, CONFIDENCE_THRESHOLD
+from app.color_detector import get_color_group
+
+CNN_MIN_SIZE = 480  
+
+
+def segment_image(image, pixels):
+    """
+    Returns a fine-grained segment_grid at the SAME resolution as `pixels`.
+    """
+
+    H, W, _ = pixels.shape
+
+    # ── 1. Run CNN on a separately upscaled copy for finer patch coverage ──
+    cnn_image = image
+    if image.width < CNN_MIN_SIZE or image.height < CNN_MIN_SIZE:
+        scale = CNN_MIN_SIZE / min(image.width, image.height)
+        new_w = int(image.width * scale)
+        new_h = int(image.height * scale)
+        cnn_image = image.resize((new_w, new_h), Image.LANCZOS)
+
+    label_grid, confidence_grid = predict_image(cnn_image)
+
+    use_cnn = bool(label_grid and label_grid[0])
+
+    if use_cnn:
+        coarse_H = len(label_grid)
+        coarse_W = len(label_grid[0])
+
+        # scale factors to map fine pixel -> coarse CNN patch
+        scale_y = coarse_H / H
+        scale_x = coarse_W / W
+
+    segment_grid = []
+
+    for y in range(H):
+        row = []
+        for x in range(W):
+
+            pixel = pixels[y, x]
+            color = get_color_group(pixel)
+
+            if use_cnn:
+                coarse_y = min(int(y * scale_y), coarse_H - 1)
+                coarse_x = min(int(x * scale_x), coarse_W - 1)
+
+                label = label_grid[coarse_y][coarse_x]
+                confidence = confidence_grid[coarse_y][coarse_x]
+
+                if confidence >= CONFIDENCE_THRESHOLD and label != "background":
+                    row.append(("segment", label, pixel))
+                else:
+                    row.append(("color", color, pixel))
+            else:
+                row.append(("color", color, pixel))
+
+        segment_grid.append(row)
+
+    return segment_grid

ml/train.py

@@ -0,0 +1,145 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, random_split
+
+from ml.dataset import build_dataset, get_transform
+from ml.model import PatchCNN
+
+
+# ── Config ──────────────────────────────────────────
+BATCH_SIZE = 64
+EPOCHS = 15
+LEARNING_RATE = 0.001
+VAL_SPLIT = 0.2
+CHECKPOINT_PATH = "checkpoints/model.pth"
+ADE20K_ROOT = "data/raw/ade20k"
+# ────────────────────────────────────────────────────
+
+
+def train():
+
+    # device — use GPU if available, otherwise CPU
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Training on: {device}")
+
+    # ── Data ──
+    print("\nLoading dataset...")
+    dataset = build_dataset(ADE20K_ROOT, transform=get_transform())
+
+    val_size = int(len(dataset) * VAL_SPLIT)
+    train_size = len(dataset) - val_size
+
+    train_set, val_set = random_split(dataset, [train_size, val_size])
+
+    train_loader = DataLoader(
+        train_set,
+        batch_size=BATCH_SIZE,
+        shuffle=True,
+        num_workers=0          # keep 0 on Windows to avoid multiprocessing issues
+    )
+
+    val_loader = DataLoader(
+        val_set,
+        batch_size=BATCH_SIZE,
+        shuffle=False,
+        num_workers=0
+    )
+
+    print(f"Train samples: {train_size}")
+    print(f"Val samples:   {val_size}")
+
+    # ── Model ──
+    model = PatchCNN(num_classes=7).to(device)
+
+    # ── Loss ──
+    # CrossEntropyLoss = softmax + negative log likelihood
+    # standard for multi-class classification
+    criterion = nn.CrossEntropyLoss()
+
+    # ── Optimizer ──
+    # Adam — adaptive learning rate, works well out of the box
+    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
+
+    # ── LR Scheduler ──
+    # reduces learning rate by 0.5 if val loss doesn't improve for 3 epochs
+    # prevents overshooting the optimal weights late in training
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer,
+        patience=3,
+        factor=0.5
+    )
+
+    # ── Training Loop ──
+    best_val_acc = 0.0
+
+    for epoch in range(EPOCHS):
+
+        # ── Train phase ──
+        model.train()
+        train_loss = 0.0
+        train_correct = 0
+
+        for batch_idx, (patches, labels) in enumerate(train_loader):
+
+            patches = patches.to(device)
+            labels = labels.to(device)
+
+            # forward pass
+            outputs = model(patches)
+            loss = criterion(outputs, labels)
+
+            # backward pass
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            train_loss += loss.item()
+
+            # count correct predictions
+            predicted = outputs.argmax(dim=1)
+            train_correct += (predicted == labels).sum().item()
+
+            # print progress every 50 batches
+            if (batch_idx + 1) % 50 == 0:
+                print(f"  Epoch {epoch+1} | Batch {batch_idx+1}/{len(train_loader)} | Loss: {loss.item():.4f}")
+
+        train_acc = train_correct / train_size
+
+        # ── Validation phase ──
+        model.eval()
+        val_loss = 0.0
+        val_correct = 0
+
+        with torch.no_grad():       # no gradients needed for validation
+            for patches, labels in val_loader:
+                patches = patches.to(device)
+                labels = labels.to(device)
+
+                outputs = model(patches)
+                loss = criterion(outputs, labels)
+
+                val_loss += loss.item()
+                predicted = outputs.argmax(dim=1)
+                val_correct += (predicted == labels).sum().item()
+
+        val_acc = val_correct / val_size
+
+        print(f"\nEpoch {epoch+1}/{EPOCHS}")
+        print(f"  Train Loss: {train_loss/len(train_loader):.4f} | Train Acc: {train_acc:.4f}")
+        print(f"  Val Loss:   {val_loss/len(val_loader):.4f} | Val Acc:   {val_acc:.4f}")
+
+        # step scheduler with val loss
+        scheduler.step(val_loss / len(val_loader))
+
+        # save best model
+        if val_acc > best_val_acc:
+            best_val_acc = val_acc
+            torch.save(model.state_dict(), CHECKPOINT_PATH)
+            print(f"  ✓ saved best model (val_acc: {val_acc:.4f})")
+
+    print(f"\nTraining complete. Best val accuracy: {best_val_acc:.4f}")
+    print(f"Model saved to {CHECKPOINT_PATH}")
+
+
+if __name__ == "__main__":
+    train()

requirements.txt

@@ -0,0 +1,10 @@
+fastapi
+uvicorn
+python-multipart
+pillow
+numpy
+torch==2.3.0+cpu
+torchvision==0.18.0+cpu
+emoji
+gradio
+--extra-index-url https://download.pytorch.org/whl/cpu