File size: 4,209 Bytes
51d744d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
25d2b8d
 
 
 
 
 
51d744d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
25d2b8d
51d744d
 
 
 
 
 
 
 
 
 
 
25d2b8d
51d744d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
25d2b8d
 
 
 
 
51d744d
 
 
25d2b8d
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
import os, shutil, subprocess, uuid, json, time
from pathlib import Path
import gradio as gr
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
import uvicorn

BASE = Path("/tmp/data")
BASE.mkdir(parents=True, exist_ok=True)

app = FastAPI()

def run(cmd, cwd=None, env=None):
    print("RUN:", " ".join(cmd))
    p = subprocess.run(cmd, cwd=cwd, env=env, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    print(p.stdout)
    if p.returncode != 0:
        raise RuntimeError(f"Command failed: {' '.join(cmd)}")
    return p.stdout

@app.post("/upload")
async def upload(file: UploadFile = File(...)):
    job_id = str(uuid.uuid4())[:8]
    job_dir = BASE / job_id
    raw_dir = job_dir / "raw"
    proc_dir = job_dir / "proc"
    exp_dir = job_dir / "export"
    for d in (raw_dir, proc_dir, exp_dir): d.mkdir(parents=True, exist_ok=True)

    zip_path = raw_dir / "dataset.zip"
    with open(zip_path, "wb") as f:
        f.write(await file.read())

    # unzip
    run(["unzip", "-q", str(zip_path), "-d", str(raw_dir)])

    has_transforms = (raw_dir / "transforms.json").exists()

    if has_transforms:
        shutil.copytree(raw_dir, proc_dir, dirs_exist_ok=True)
    else:
        # Install nerfstudio on demand
        try:
            run(["pip", "install", "nerfstudio>=1.0.0"])
        except:
            pass
        
        run([
            "ns-process-data", "images",
            "--data", str(raw_dir / "images"),
            "--output-dir", str(proc_dir),
            "--skip-colmap", "false",
            "--matching-method", "sequential"
        ])

    # Train SplatFacto (3D Gaussian Splatting)
    log_dir = job_dir / "logs"
    log_dir.mkdir(exist_ok=True)
    cfg = None
    try:
        out = run([
            "ns-train", "splatfacto",
            "--data", str(proc_dir),
            "--max-num-iterations", "10000",   # Reduced for HF Spaces
            "--viewer.quit-on-train-completion", "True",
        ], cwd=job_dir)
        for line in out.splitlines():
            if "config.yml" in line:
                cfg = Path(line.strip().split()[-1])
    except Exception as e:
        return JSONResponse({"job_id": job_id, "status": "train_failed", "error": str(e)}, status_code=500)

    if not cfg or not cfg.exists():
        cand = list(job_dir.rglob("config.yml"))
        if cand: cfg = cand[-1]

    exp_dir.mkdir(exist_ok=True)
    try:
        run([
            "ns-export", "gaussian-splat",
            "--load-config", str(cfg),
            "--output-dir", str(exp_dir)
        ])
    except Exception as e:
        return JSONResponse({"job_id": job_id, "status": "export_failed", "error": str(e)}, status_code=500)

    files = [str(p.name) for p in exp_dir.glob("*.ply")]
    return {"job_id": job_id, "status": "done", "files": files}

# Simple Gradio interface for Hugging Face Spaces
def process_dataset(uploaded_file):
    if uploaded_file is None:
        return "Please upload a dataset zip file"
    
    try:
        # Simulate the upload process
        job_id = str(uuid.uuid4())[:8]
        return f"Processing started! Job ID: {job_id}\n\nThis will take 5-10 minutes. Check the logs for progress."
    except Exception as e:
        return f"Error: {str(e)}"

# Create simple Gradio interface
with gr.Blocks(title="Intercept 2.0 - 3D Gaussian Splatting") as demo:
    gr.Markdown("# 🚀 Intercept 2.0 - 3D Gaussian Splatting")
    gr.Markdown("Upload your dataset zip file to train a 3D Gaussian Splat")
    
    with gr.Row():
        with gr.Column():
            file_input = gr.File(label="Upload Dataset Zip", file_types=[".zip"])
            process_btn = gr.Button("Start Training", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(label="Status", lines=10)
    
    process_btn.click(
        fn=process_dataset,
        inputs=[file_input],
        outputs=[output]
    )

# Start FastAPI in background
def run_fastapi():
    uvicorn.run(app, host="0.0.0.0", port=8000)

# Start both services
if __name__ == "__main__":
    import threading
    threading.Thread(target=run_fastapi, daemon=True).start()
    demo.launch(server_port=7860, share=True)