Update app.py

app.py

@@ -1,7 +1,5 @@
 import gradio as gr
-import torch
-import numpy as np
-import json, os, sys, glob, shutil
+import json, os, sys, glob, shutil, subprocess
 from PIL import Image
 
 #os.system("git clone https://github.com/sunset1995/HorizonNet.git /tmp/HorizonNet")
@@ -12,6 +10,15 @@ from huggingface_hub import hf_hub_download
 CKPT = hf_hub_download(repo_id="gum-tech/horizonnet-resnet50-rnn", filename="resnet50_rnn__st3d.pth")
 print("Model checkpoint at:", CKPT)
 
+def run(cmd, cwd="/tmp/HorizonNet"):
+    result = subprocess.run(
+        cmd, cwd=cwd, shell=False,
+        stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+    )
+    log = result.stdout.decode("utf-8", errors="replace")
+    print(f"=== {cmd[0]} (exit {result.returncode}) ===\n{log}\n===")
+    return result.returncode, log
+
 def predict(image):
     for d in ["/tmp/hn_input", "/tmp/hn_pre", "/tmp/hn_out"]:
         shutil.rmtree(d, ignore_errors=True)
@@ -20,36 +27,32 @@ def predict(image):
     img_path = "/tmp/hn_input/room.png"
     image.convert("RGB").resize((1024, 512)).save(img_path)
 
-    os.system(
-        "cd /tmp/HorizonNet && python preprocess.py"
-        " --img_glob " + img_path +
-        " --output_dir /tmp/hn_pre"
-        " > /tmp/hn_pre/preprocess_log.txt 2>&1"
-    )
-    log = open("/tmp/hn_pre/preprocess_log.txt").read() if os.path.exists("/tmp/hn_pre/preprocess_log.txt") else "no log"
-    print("=== preprocess log ===\n", log)
+    # Step 1: preprocess
+    code, log = run([
+        "python", "preprocess.py",
+        "--img_glob", img_path,
+        "--output_dir", "/tmp/hn_pre"
+    ])
 
     aligned = glob.glob("/tmp/hn_pre/*_aligned_rgb.png")
     if not aligned:
-        return json.dumps({"error": "Preprocessing failed", "log": log})
+        return json.dumps({"error": "Preprocessing failed", "code": code, "log": log[-1000:]})
 
     aligned_path = aligned[0]
     print("Aligned image:", aligned_path)
 
-    os.system(
-        "cd /tmp/HorizonNet && python inference.py"
-        " --pth " + CKPT +
-        " --img_glob " + aligned_path +
-        " --output_dir /tmp/hn_out"
-        " --no_cuda"
-        " > /tmp/hn_out/inference_log.txt 2>&1"
-    )
-    inf_log = open("/tmp/hn_out/inference_log.txt").read() if os.path.exists("/tmp/hn_out/inference_log.txt") else "no log"
-    print("=== inference log ===\n", inf_log)
+    # Step 2: inference
+    code, log = run([
+        "python", "inference.py",
+        "--pth", CKPT,
+        "--img_glob", aligned_path,
+        "--output_dir", "/tmp/hn_out",
+        "--no_cuda"
+    ])
 
     out_jsons = glob.glob("/tmp/hn_out/*.json")
     if not out_jsons:
-        return json.dumps({"error": "Inference failed", "log": inf_log})
+        return json.dumps({"error": "Inference failed", "code": code, "log": log[-1000:]})
 
     with open(out_jsons[0]) as f:
         result = json.load(f)