Update app.py

app.py

@@ -16,6 +16,7 @@ from torchvision.transforms import Compose
 from tqdm import tqdm
 import TEED.main as teed
 from TEED.main import parse_args
+import gradio as gr 
 
 from depthAnything.depth_anything.dpt import DepthAnything
 from depthAnything.depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet
@@ -342,14 +343,17 @@ def process_images(input_folder='./output/merged_imgs'):
         except Exception as e:
             print(f"Error processing file {image_file}: {e}")
 
-def process_line(img_path='./input', outdir='./output'):
+def process_line(img, outdir='./output'):
+    img_path = "./input"
+    cv2.imwrite(os.path.join(img_path, "input_image.png"), img)  # 保存上传的图像
     depth_anything(img_path, os.path.join(outdir, "depth_anything"))
     teed_imgs(img_path, os.path.join(outdir, "teed_imgs"), [1, 7, 2])
     teed_imgs(os.path.join(outdir, "depth_anything"), os.path.join(outdir, "dp_teed_imgs"), [0, 7, 2])
     merge_images_in_2_folder(os.path.join(outdir, "teed_imgs"), os.path.join(outdir, "dp_teed_imgs"), os.path.join(outdir, "merged_imgs"), '_depth', 1, 'multiply', [[2, 0], [2, 1]], [1, 0])
-    process_images(os.path.join(outdir, "merged_imgs"))  # 处理merged_imgs文件夹中的图片
-
-
+    
+    # 返回处理后的最后一张图像
+    result_img_path = os.path.join(outdir, "merged_imgs", "your_final_image_name.png")  # 修改为你的输出图像路径
+    return Image.open(result_img_path)
 
 
 if __name__ == '__main__':
@@ -365,3 +369,15 @@ if __name__ == '__main__':
     img_path = "input"
     outdir = "output"
     process_line(img_path,outdir)
+
+
+if __name__ == '__main__':
+    # 创建 Gradio 接口
+    interface = gr.Interface(fn=process_line,
+                             inputs=gr.Image(type="numpy", label="Upload Image"),
+                             outputs=gr.Image(label="Output Image"),
+                             title="Image Processing",
+                             description="Upload an image to process it using depth-anything and TEED.")
+
+    # 启动 Gradio 接口
+    interface.launch()