Upload folder using huggingface_hub

predict.py

@@ -106,6 +106,18 @@ def load_model(model_dir: str):
     return ens
 
 
+def _unsharp_mask(frames_tensor, alpha=0.15):
+    """Apply unsharp mask sharpening. frames_tensor: [1, N, 3, 64, 64]"""
+    import torch.nn.functional as F
+    kernel = torch.tensor([[1, 2, 1], [2, 4, 2], [1, 2, 1]], dtype=frames_tensor.dtype, device=frames_tensor.device).float() / 16.0
+    kernel = kernel.unsqueeze(0).unsqueeze(0).expand(3, 1, 3, 3)
+    b, n, c, h, w = frames_tensor.shape
+    x = frames_tensor.reshape(b * n, c, h, w)
+    blurred = F.conv2d(x, kernel, padding=1, groups=3)
+    sharpened = x + alpha * (x - blurred)
+    return torch.clamp(sharpened.reshape(b, n, c, h, w), 0, 1)
+
+
 def _predict_8frames_direct(model, context_tensor, last_tensor):
     output = model(context_tensor)
     residuals = output.reshape(1, PRED_FRAMES, 3, 64, 64)
@@ -236,6 +248,8 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
                 direct_weight = 1.0 - ar_weight
                 predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
 
+            predicted = _unsharp_mask(predicted, alpha=0.15)
+
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
@@ -261,19 +275,13 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             context_tensor = torch.from_numpy(context).to(DEVICE)
             last_tensor = torch.from_numpy(last_frame_t).to(DEVICE)
 
+            predicted_orig = _predict_8frames_direct(ens.models["pole_position"], context_tensor, last_tensor)
             context_flipped = torch.flip(context_tensor, dims=[3])
             last_flipped = torch.flip(last_tensor, dims=[3])
-
-            # 3-run noise diversity + TTA for PP
-            all_pp_runs = []
-            for noise_std in [0.0, 0.5/255.0, 1.0/255.0]:
-                ctx_in = context_tensor if noise_std == 0 else torch.clamp(context_tensor + torch.randn_like(context_tensor) * noise_std, 0, 1)
-                ctx_flip_in = context_flipped if noise_std == 0 else torch.clamp(context_flipped + torch.randn_like(context_flipped) * noise_std, 0, 1)
-                pred_orig = _predict_8frames_direct(ens.models["pole_position"], ctx_in, last_tensor)
-                pred_flipped = _predict_8frames_direct(ens.models["pole_position"], ctx_flip_in, last_flipped)
-                pred_flipped = torch.flip(pred_flipped, dims=[4])
-                all_pp_runs.append((pred_orig + pred_flipped) / 2.0)
-            predicted = sum(all_pp_runs) / len(all_pp_runs)
+            predicted_flipped = _predict_8frames_direct(ens.models["pole_position"], context_flipped, last_flipped)
+            predicted_flipped = torch.flip(predicted_flipped, dims=[4])
+            predicted = (predicted_orig + predicted_flipped) / 2.0
+            predicted = _unsharp_mask(predicted, alpha=0.15)
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []