Update gradio_depth_pred.py

gradio_depth_pred.py

@@ -6,7 +6,7 @@ import torchvision.transforms as transforms
 import torch
 import spaces
 
-resolution = 2
+resolution = 4
 base_area = resolution * 480 * 640
 flip_test = True
 
@@ -27,72 +27,16 @@ def predict_depth(model, image):
     new_w, new_h = int(scale * frame_width), int(scale * frame_height)
     frame = cv2.resize(frame, (new_w, new_h))
     
-    # ---------------------------------------
-    # Test-Time Augmentations (TTA) + flip-test for each augmentation
-    # ---------------------------------------
-    
-    frame_tensor = transforms.ToTensor()(frame).unsqueeze(0).to(DEVICE)
-    
-    # ---- Base augmentations (without flip) ----
-    def apply_resize(x):
-        b, c, h, w = x.shape
-        scale = torch.empty(1).uniform_(0.9, 1.1).item()
-        nh, nw = int(h * scale), int(w * scale)
-        x2 = torch.nn.functional.interpolate(x, (nh, nw), mode="bilinear", align_corners=False)
-        x2 = torch.nn.functional.interpolate(x2, (h, w), mode="bilinear", align_corners=False)
-        return x2
-    
-    brightness_jitter = transforms.ColorJitter(brightness=0.15)
-    def apply_brightness(x):
-        img = transforms.ToPILImage()(x[0].cpu())
-        img = brightness_jitter(img)
-        return transforms.ToTensor()(img).unsqueeze(0).to(x.device)
-    
-    tta_base = [
-        ("orig", lambda x: x),
-        ("resize", apply_resize),
-        ("bright", apply_brightness),
-    ]
-    
-    # ---------------------------------------
-    # Build augmented batch (A(x) and A(x_flipped))
-    # ---------------------------------------
-    augmented_frames = []
-    reverse_fns = []
-    
-    for name, aug_fn in tta_base:
-        # A(x)
-        ax = aug_fn(frame_tensor)
-        augmented_frames.append(ax)
-        reverse_fns.append(lambda y: y)       # no unflip needed
-    
-        if True:  # flip_test
-            # A(x_flip)
-            axf = aug_fn(frame_tensor.flip(-1))
-            augmented_frames.append(axf)
-            reverse_fns.append(lambda y: y.flip(-1))  # unflip prediction
-    
-    batch = torch.cat(augmented_frames, dim=0)   # [N_aug*2, 3, H, W]
-    
-    # ---------------------------------------
-    # One forward pass
-    # ---------------------------------------
+    frame = transforms.ToTensor()(frame).unsqueeze(0)
+    if flip_test:
+        frame = torch.cat([frame, frame.flip(-1)])
+    frame = frame.to(DEVICE)
     model.to(DEVICE)
-    
     with torch.no_grad():
-        pred_batch = model(batch)   # [N_aug*2, 1, H, W]
-    
-    # ---------------------------------------
-    # Reverse and average predictions
-    # ---------------------------------------
-    corrected = []
-    for i, reverse in enumerate(reverse_fns):
-        corrected.append(reverse(pred_batch[i:i+1]))
-    
-    depth = torch.stack(corrected).mean(dim=0)
-    depth = depth[0, 0].cpu().numpy()
-    return depth
-
+        depth = model(frame)
+    if flip_test:
+        depth = ((depth[0] + depth[1].flip(-1))/2).unsqueeze(0)
+    return depth.detach().cpu().numpy()[0, 0]
 
 def create_demo(model, scene):
     gr.Markdown("### Depth Prediction demo")