Update app.py

app.py

@@ -8,18 +8,13 @@ import shutil
 import os
 import teed
 
-
 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等
-# 这些函数无需修改
-
-# 修改depth_anything函数，直接处理传入的图片
+# 深度处理函数
 def depth_anything_image(image, encoder='vitl', pred_only=True, grayscale=True):
     DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-
     model_configs = {
         'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
         'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
@@ -37,11 +32,8 @@ def depth_anything_image(image, encoder='vitl', pred_only=True, grayscale=True):
         PrepareForNet(),
     ])
 
-    raw_image = np.array(image.convert('RGB'))
-    raw_image = raw_image[:, :, ::-1].copy()  # Convert RGB to BGR for OpenCV
-
+    raw_image = np.array(image.convert('RGB'))[:, :, ::-1].copy()  # RGB to BGR
     h, w = raw_image.shape[:2]
-
     image_tensor = transform({'image': raw_image / 255.0})['image']
     image_tensor = torch.from_numpy(image_tensor).unsqueeze(0).to(DEVICE)
 
@@ -52,64 +44,39 @@ def depth_anything_image(image, encoder='vitl', pred_only=True, grayscale=True):
     depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
     depth = depth.cpu().numpy().astype(np.uint8)
 
-    if grayscale:
-        depth = np.repeat(depth[..., np.newaxis], 3, axis=-1)
-    else:
-        depth = cv2.applyColorMap(depth, cv2.COLORMAP_INFERNO)
-
-    if pred_only:
-        return depth
-    else:
-        combined_results = np.hstack([raw_image, depth])
-        return combined_results
+    return np.repeat(depth[..., np.newaxis], 3, axis=-1) if grayscale else cv2.applyColorMap(depth, cv2.COLORMAP_INFERNO)
 
-# TEED function modified to process a single image
+# TEED 图像处理函数
 def teed_process_image(image):
     os.makedirs('./output/teed_imgs', exist_ok=True)
-    os.makedirs('teed_tmp', exist_ok=True)
+    os.makedirs('./teed_tmp', exist_ok=True)
 
-    # 将图像保存为临时文件
     temp_image_path = './teed_tmp/temp_image.png'
     cv2.imwrite(temp_image_path, np.array(image))
 
-    # 设置 TEED 处理所需的参数
-    args, train_info = parse_args(is_testing=True, pl_opt_dir='./output/teed_imgs')
+    args = parse_args(is_testing=True, pl_opt_dir='./output/teed_imgs')
     args.input_val_dir = './teed_tmp'  # 临时目录
+    args.output_dir = './output/teed_imgs'  # 输出目录
 
-    # 调用 TEED 主函数进行处理
-    teed.main(args, train_info)
-
-    # 清理临时文件
-    shutil.rmtree('teed_tmp')
+    # 确保 TEED 处理的代码正常
+    if hasattr(teed, 'main'):
+        teed.main(args)
+    else:
+        print("TEED module does not have a main function.")
 
+    shutil.rmtree('./teed_tmp')
     return cv2.imread(os.path.join('./output/teed_imgs', 'processed_image.png'))
 
-
-    return processed_image
-
-# 修改处理流程，处理单个上传的图片
+# 处理单个图像
 def process_single_image(image):
-    # 深度处理
     depth_result = depth_anything_image(image, 'vitl')
-
-    # TEED 边缘检测
     teed_result = teed_process_image(image)
-
-    # 将两个结果叠加合并（例如 Multiply）
     merged_result = multiply_blend(depth_result, teed_result)
-
     return merged_result
 
-# Gradio界面处理函数
+# 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)
 
 # Gradio 界面