app.py

import gradio as gr
import torch
from transformers import DPTForDepthEstimation, DPTImageProcessor
from PIL import Image
import numpy as np
import trimesh
from nbtschematic import SchematicFile
from scipy.spatial import cKDTree
import os
import shutil
import zipfile

# --- NEW IMPORTS for .mcworld export ---

# Initialize to a safe default state
AMULET_AVAILABLE = False
PALETTE_AMULET_BLOCKS = []

try:
    import amulet
    import amulet.api.world
    from amulet.api.selection import SelectionGroup, SelectionBox
    from amulet.api.structure import Structure
    from amulet.api.block import Block
    
    # If imports succeed, set the flag to True
    AMULET_AVAILABLE = True
    print("Amulet-core loaded successfully.")

except ImportError:
    # If imports fail, just print a warning. The flag remains False.
    print("Amulet-core not found. .mcworld export will be disabled.")


# --- 1. SETUP ---

print("Initializing... Loading AI model.")
processor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas")
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas")
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
print(f"Model loaded to device: {device}")

MINECRAFT_PALETTE_DATA = {
    "Stone": (("minecraft", "stone"), (1, 0), (128, 128, 128)),
    "Cobblestone": (("minecraft", "cobblestone"), (4, 0), (125, 125, 125)),
    "Gravel": (("minecraft", "gravel"), (13, 0), (136, 132, 131)),
    "Andesite": (("minecraft", "stone", {"stone_type": "andesite"}), (1, 5), (136, 136, 136)),
    "Iron Block": (("minecraft", "iron_block"), (42, 0), (218, 218, 218)),
    "Clay Block": (("minecraft", "clay"), (82, 0), (164, 172, 183)),
    "Light Gray Wool": (("minecraft", "wool", {"color": "light_gray"}), (35, 8), (142, 142, 134)),
    "Gray Wool": (("minecraft", "wool", {"color": "gray"}), (35, 7), (65, 68, 72)),
    "Black Wool": (("minecraft", "wool", {"color": "black"}), (35, 15), (21, 21, 26)),
    "White Wool": (("minecraft", "wool", {"color": "white"}), (35, 0), (234, 236, 237)),
    "Oak Planks": (("minecraft", "planks", {"wood_type": "oak"}), (5, 0), (157, 128, 79)),
    "Spruce Planks": (("minecraft", "planks", {"wood_type": "spruce"}), (5, 1), (101, 75, 43)),
    "Birch Planks": (("minecraft", "planks", {"wood_type": "birch"}), (5, 2), (193, 174, 114)),
    "Dark Oak Planks": (("minecraft", "planks", {"wood_type": "dark_oak"}), (5, 5), (61, 42, 21)),
    "Dirt": (("minecraft", "dirt"), (3, 0), (134, 96, 67)),
    "Red Wool": (("minecraft", "wool", {"color": "red"}), (35, 14), (168, 50, 51)),
    "Sand": (("minecraft", "sand"), (12, 0), (218, 211, 160)),
    "Grass Block": (("minecraft", "grass_block"), (2, 0), (123, 182, 74)),
    "Blue Wool": (("minecraft", "wool", {"color": "blue"}), (35, 11), (53, 57, 157)),
    "Light Blue Wool": (("minecraft", "wool", {"color": "light_blue"}), (35, 3), (105, 158, 210)),
    "Obsidian": (("minecraft", "obsidian"), (49, 0), (22, 18, 29)),
}

# --- This logic is now conditional ---
PALETTE_VALUES = list(MINECRAFT_PALETTE_DATA.values())
if AMULET_AVAILABLE:
    # Only create the Amulet Block list if the library was imported successfully
    PALETTE_AMULET_BLOCKS = [Block(v[0][0], v[0][1], v[0][2] if len(v[0]) > 2 else {}) for v in PALETTE_VALUES]

PALETTE_SCHEMATIC_BLOCKS = np.array([v[1] for v in PALETTE_VALUES])
PALETTE_COLORS = np.array([v[2] for v in PALETTE_VALUES])
COLOR_TREE = cKDTree(PALETTE_COLORS)
print("Palette initialized.")


# --- CORE & EXPORT FUNCTIONS (Unchanged from previous attempt) ---

def predict_depth(image):
    inputs = processor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
        predicted_depth = outputs.predicted_depth
    prediction = torch.nn.functional.interpolate(
        predicted_depth.unsqueeze(1), size=image.size[::-1], mode="bicubic", align_corners=False
    )
    output = prediction.squeeze().cpu().numpy()
    return (output - output.min()) / (output.max() - output.min())

def create_point_cloud_from_depth(depth_map, image, max_voxels=128):
    h, w = depth_map.shape
    scale_factor = np.sqrt((max_voxels**2) / (w * h)) if w*h > 0 else 0
    new_w, new_h = int(w * scale_factor), int(h * scale_factor)
    if new_w == 0 or new_h == 0: return np.array([]), np.array([])
    small_depth = np.array(Image.fromarray(depth_map).resize((new_w, new_h)))
    small_image = image.resize((new_w, new_h))
    colors = np.array(small_image).reshape(-1, 3)
    y, x = np.mgrid[:new_h, :new_w]
    lon = (x / new_w) * 2 * np.pi - np.pi
    lat = -(y / new_h) * np.pi + np.pi / 2
    r = small_depth * (max_voxels / 2)
    px = r * np.cos(lat) * np.cos(lon)
    py = r * np.cos(lat) * np.sin(lon)
    pz = r * np.sin(lat)
    points = np.stack([px, py, pz], axis=-1).reshape(-1, 3)
    points -= points.mean(axis=0)
    return points, colors

def voxelize(points, colors, pitch=1.0):
    if len(points) == 0: return np.array([]), np.array([])
    scaled_points = points / pitch
    cloud = trimesh.points.PointCloud(scaled_points, colors=colors)
    voxel_grid = trimesh.voxel.VoxelGrid(cloud)
    if not isinstance(voxel_grid, trimesh.voxel.VoxelGrid): return np.array([]), np.array([])
    voxel_coords, voxel_colors = voxel_grid.matrix_to_frame()
    if len(voxel_colors) == 0: return np.array([]), np.array([])
    _, indices = COLOR_TREE.query(voxel_colors[:, :3])
    return voxel_coords, indices

def create_schematic(voxel_coords, palette_indices, filename):
    if len(voxel_coords) == 0:
        sf = SchematicFile(shape=(1, 1, 1)); sf.save(filename)
        return filename
    voxel_block_data = PALETTE_SCHEMATIC_BLOCKS[palette_indices]
    min_c = voxel_coords.min(axis=0); max_c = voxel_coords.max(axis=0)
    dims = (max_c - min_c + 1).astype(int)
    sf = SchematicFile(shape=(dims[2], dims[1], dims[0]))
    for i, coord in enumerate(voxel_coords):
        x, y, z = (coord - min_c).astype(int)
        block_id, block_data = voxel_block_data[i]
        if 0 <= z < sf.shape[0] and 0 <= y < sf.shape[1] and 0 <= x < sf.shape[2]:
            sf.Blocks[z, y, x] = block_id
            sf.Data[z, y, x] = block_data
    sf.save(filename)
    return filename

def create_mcworld(voxel_coords, palette_indices, filename):
    if not AMULET_AVAILABLE:
        raise gr.Error("Amulet library failed to load. Cannot create .mcworld files.")
    TEMP_WORLD_DIR = "./temp_mc_world"
    if os.path.exists(TEMP_WORLD_DIR): shutil.rmtree(TEMP_WORLD_DIR)
    world = amulet.load_format(TEMP_WORLD_DIR)
    world.create_and_open("bedrock", (0, 255))
    if len(voxel_coords) > 0:
        min_c = voxel_coords.min(axis=0)
        selection_box = SelectionBox(min_c, (voxel_coords.max(axis=0) - min_c + 1))
        structure_array = np.zeros(selection_box.shape, dtype=object)
        for i, coord in enumerate(voxel_coords):
            x, y, z = (coord - min_c).astype(int)
            block = PALETTE_AMULET_BLOCKS[palette_indices[i]]
            structure_array[x, y, z] = (block, None)
        structure = Structure(SelectionGroup(selection_box), structure_array, {})
        world.put_structure(structure, world.dimensions[0])
    world.save(); world.close()
    with zipfile.ZipFile(filename, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(TEMP_WORLD_DIR):
            for file in files:
                filepath = os.path.join(root, file)
                zipf.write(filepath, filepath.replace(TEMP_WORLD_DIR, ''))
    shutil.rmtree(TEMP_WORLD_DIR)
    return filename


# --- MAIN LOGIC & UI ---

def panorama_to_minecraft(pano_image, max_dim, output_format):
    if pano_image is None: raise gr.Error("Please upload a 360° panorama image.")
    if output_format == ".mcworld" and not AMULET_AVAILABLE:
        raise gr.Error("Amulet library is not available in this environment. Please choose .schematic instead.")
    print("Step 1/4: Predicting Depth...")
    depth_map = predict_depth(pano_image)
    print("Step 2/4: Creating 3D Point Cloud...")
    points, colors = create_point_cloud_from_depth(depth_map, pano_image, max_voxels=max_dim)
    print("Step 3/4: Voxelizing Scene...")
    voxel_coords, palette_indices = voxelize(points, colors, pitch=1.0)
    print(f"Step 4/4: Creating {output_format} file...")
    if not os.path.exists("outputs"): os.makedirs("outputs")
    if output_format == ".mcworld":
        output_filename = os.path.join("outputs", "PanoCraft.mcworld")
        schematic_path = create_mcworld(voxel_coords, palette_indices, output_filename)
    else:
        output_filename = os.path.join("outputs", "PanoCraft.schematic")
        schematic_path = create_schematic(voxel_coords, palette_indices, output_filename)
    print("Processing complete.")
    return schematic_path

with gr.Blocks(title="PanoCraft") as demo:
    gr.Markdown("# 360° Panorama to Minecraft Converter")
    gr.Markdown("Upload a 360° photo to create a Minecraft model. Export as a `.schematic` for editors like Amulet, or as a `.mcworld` to import directly into Bedrock Edition (PE, Windows, etc).")
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="pil", label="Upload 360° Panorama Image")
            max_dim_slider = gr.Slider(minimum=32, maximum=256, value=128, step=16, label="Detail Level", info="Higher values create more detailed models.")
            output_format_radio = gr.Radio(
                [".schematic", ".mcworld"], value=".schematic", label="Output Format", info=".mcworld is for direct Bedrock import."
            )
            submit_btn = gr.Button("Generate Minecraft File", variant="primary")
        with gr.Column(scale=1):
            file_output = gr.File(label="Download File")
    submit_btn.click(
        fn=panorama_to_minecraft,
        inputs=[image_input, max_dim_slider, output_format_radio],
        outputs=file_output
    )
    gr.Markdown("### Example")
    gr.Examples(
        examples=[[os.path.join(os.path.dirname(__file__), "example_pano.jpg")]],
        inputs=[image_input],
        label="Click to try"
    )

demo.launch()