Update app.py

app.py

@@ -7,8 +7,10 @@ from torchvision.transforms import Compose
 import shutil
 import os
 
+
 from depthAnything.depth_anything.dpt import DepthAnything
 from depthAnything.depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet
+from TEED.main import parse_args
 
 # 保留原有图像处理函数如multiply_blend、screen_blend、erosion等
 # 这些函数无需修改
@@ -61,24 +63,26 @@ def depth_anything_image(image, encoder='vitl', pred_only=True, grayscale=True):
         return combined_results
 
 # TEED function modified to process a single image
-def teed_process_image(image, gaussianBlur=[0, 3, 0]):
-    temp_image_path = './teed_temp_image.png'
+def teed_process_image(image):
+    os.makedirs('./output/teed_imgs', exist_ok=True)
     os.makedirs('teed_tmp', exist_ok=True)
 
-    img = np.array(image.convert('RGB'))[:, :, ::-1].copy()  # Convert RGB to BGR for OpenCV
-
-    if gaussianBlur[0] != 0:
-        img = cv2.GaussianBlur(img, (gaussianBlur[1], gaussianBlur[1]), gaussianBlur[2])
+    # 将图像保存为临时文件
+    temp_image_path = './teed_tmp/temp_image.png'
+    cv2.imwrite(temp_image_path, np.array(image))
 
-    cv2.imwrite(temp_image_path, img)
-
-    # Run TEED on the image
+    # 设置 TEED 处理所需的参数
     args, train_info = parse_args(is_testing=True, pl_opt_dir='./output/teed_imgs')
-    args.input_val_dir = 'teed_tmp'  # Temporary directory for single image processing
+    args.input_val_dir = './teed_tmp'  # 临时目录
+
+    # 调用 TEED 主函数进行处理
     teed.main(args, train_info)
 
-    processed_image = cv2.imread(temp_image_path)  # Load the processed image
-    shutil.rmtree('teed_tmp')  # Clean up the temporary directory
+    # 清理临时文件
+    shutil.rmtree('teed_tmp')
+
+    return cv2.imread(os.path.join('./output/teed_imgs', 'processed_image.png'))
+
 
     return processed_image
 
@@ -97,8 +101,15 @@ def process_single_image(image):
 
 # Gradio界面处理函数
 def gradio_process_line(img):
+    # 保存上传的图像到临时路径
+    img_path = './temp_input.png'
+    img.save(img_path)
+
+    # 处理图像
     processed_image = process_single_image(img)
-    return Image.fromarray(processed_image[:, :, ::-1])  # Convert BGR back to RGB for PIL
+
+    # 返回处理后的图像
+    return Image.fromarray(processed_image)
 
 # Gradio 界面
 import gradio as gr