app.py

import os
import sys
import argparse
import numpy as np
import cv2
from PIL import Image
import tkinter as tk
from tkinter import filedialog

CONFIG = {
    "INPUT": "car.jpg",
    "USE_MIDAS": False,
    "MODEL_TYPE": "MiDaS_small",
    "DEPTH_SCALE": 1.0,
    "SAMPLE_STEP": 3,
    "POINT_SIZE": 2.0,
}

def choose_image_dialog():
    root = tk.Tk()
    root.withdraw()
    filetypes = [("Image files", "*.jpg *.jpeg *.png *.bmp *.tiff"), ("All files", "*.*")]
    path = filedialog.askopenfilename(title="Select image", filetypes=filetypes)
    root.update()
    root.destroy()
    return path

def read_rgb(path: str) -> np.ndarray:
    img_bgr = cv2.imread(path, cv2.IMREAD_COLOR)
    if img_bgr is None:
        raise FileNotFoundError(f"Cannot read image: {path}")
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    return img_rgb

def estimate_depth_simple(rgb: np.ndarray) -> np.ndarray:
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY).astype(np.float32) / 255.0
    gray = cv2.GaussianBlur(gray, (9, 9), 0)
    depth = 1.0 - gray
    mn, mx = depth.min(), depth.max()
    if mx - mn > 1e-6:
        depth = (depth - mn) / (mx - mn)
    else:
        depth = np.zeros_like(depth)
    return depth

def estimate_depth_midas(rgb: np.ndarray, model_type: str = "MiDaS_small", device_preference: str = None) -> np.ndarray:
    try:
        import torch
    except Exception as e:
        raise RuntimeError("MiDaS requires 'torch'. Install it with 'pip install torch torchvision'") from e

    device = device_preference or ("cuda" if torch.cuda.is_available() else "cpu")
    midas = torch.hub.load("intel-isl/MiDaS", model_type)
    midas.to(device).eval()
    midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
    if model_type == "MiDaS_small":
        transform = midas_transforms.small_transform
    else:
        transform = midas_transforms.dpt_transform

    pil = Image.fromarray(rgb)
    input_batch = transform(pil).unsqueeze(0).to(device)
    with torch.no_grad():
        prediction = midas(input_batch)
        prediction = torch.nn.functional.interpolate(
            prediction,
            size=rgb.shape[:2],
            mode="bicubic",
            align_corners=False,
        ).squeeze()
    depth = prediction.cpu().numpy()
    mn, mx = depth.min(), depth.max()
    if mx - mn > 1e-6:
        depth = (depth - mn) / (mx - mn)
    else:
        depth = np.zeros_like(depth)
    return depth

def create_point_cloud(rgb: np.ndarray, depth: np.ndarray, sample_step: int = 3, depth_scale: float = 1.0):
    h, w = depth.shape
    ys = np.arange(0, h, sample_step)
    xs = np.arange(0, w, sample_step)
    xv, yv = np.meshgrid(xs, ys)
    xv = xv.flatten()
    yv = yv.flatten()
    z = depth[yv, xv].astype(np.float32) * float(depth_scale)

    px_scale = 1.0 / max(w, h)
    cx = w / 2.0
    cy = h / 2.0
    x = (xv - cx) * px_scale
    y = (yv - cy) * px_scale * (-1)
    points = np.stack([x, y, z], axis=1)
    colors = rgb[yv, xv].astype(np.float32) / 255.0
    return points, colors

def visualize_points(points: np.ndarray, colors: np.ndarray, point_size: float = 2.0):
    try:
        import open3d as o3d
    except Exception as e:
        raise RuntimeError("Open3D is required. Install with 'pip install open3d'") from e

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    pcd.colors = o3d.utility.Vector3dVector(colors)

    vis = o3d.visualization.Visualizer()
    vis.create_window(window_name="3D Viewer (rotate with mouse)", width=1200, height=800)
    vis.add_geometry(pcd)
    opt = vis.get_render_option()
    opt.point_size = float(point_size)
    opt.background_color = np.asarray([0.0, 0.0, 0.0])
    axis = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.2, origin=[0,0,0])
    vis.add_geometry(axis)
    ctr = vis.get_view_control()
    ctr.set_zoom(0.8)
    print("Viewer opened. Use mouse to rotate, middle mouse to zoom/pan.")
    vis.run()
    vis.destroy_window()

def main(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", "-i", help="Input image path")
    parser.add_argument("--use_midas", action="store_true")
    parser.add_argument("--model_type", default=CONFIG["MODEL_TYPE"])
    parser.add_argument("--depth_scale", type=float, default=CONFIG["DEPTH_SCALE"])
    parser.add_argument("--sample_step", type=int, default=CONFIG["SAMPLE_STEP"])
    parser.add_argument("--point_size", type=float, default=CONFIG["POINT_SIZE"])
    args = parser.parse_args(argv)

    cfg = CONFIG.copy()
    if args.input:
        cfg["INPUT"] = args.input
    cfg["USE_MIDAS"] = args.use_midas
    cfg["MODEL_TYPE"] = args.model_type
    cfg["DEPTH_SCALE"] = args.depth_scale
    cfg["SAMPLE_STEP"] = args.sample_step
    cfg["POINT_SIZE"] = args.point_size

    input_path = cfg.get("INPUT")
    if input_path and not os.path.exists(input_path):
        print(f"{input_path} not found. Opening file picker.")
        input_path = None

    if not input_path:
        try:
            input_path = choose_image_dialog()
        except Exception:
            print("Could not open file picker. Exiting.")
            return
        if not input_path:
            print("No image selected. Exiting.")
            return

    print("Using image:", input_path)
    rgb = read_rgb(input_path)

    if cfg.get("USE_MIDAS"):
        try:
            depth = estimate_depth_midas(rgb, model_type=cfg.get("MODEL_TYPE", "MiDaS_small"))
        except Exception as e:
            print("MiDaS failed, falling back to simple depth. Error:", e)
            depth = estimate_depth_simple(rgb)
    else:
        depth = estimate_depth_simple(rgb)

    points, colors = create_point_cloud(
        rgb,
        depth,
        sample_step=int(cfg.get("SAMPLE_STEP", 3)),
        depth_scale=float(cfg.get("DEPTH_SCALE", 1.0))
    )

    print(f"Point count: {len(points)} (sample_step = {cfg.get('SAMPLE_STEP')})")
    visualize_points(points, colors, point_size=cfg.get("POINT_SIZE", 2.0))

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print("Error:", e)
        sys.exit(1)