Update inference.py

inference.py

@@ -90,7 +90,7 @@ class DTDPredictor:
         The model was trained on quantized coefficients, so results may vary.
 
         Args:
-            image_array: RGB image as numpy array
+            image_array: RGB image as numpy array (must be 8x8 aligned)
             quality: JPEG quality (used for QT approximation)
 
         Returns:
@@ -101,11 +101,12 @@ class DTDPredictor:
         im_ycbcr = cv2.cvtColor(image_array, cv2.COLOR_RGB2YCrCb)
         y_channel = im_ycbcr[:, :, 0].astype(np.float32) - 128  # Center around 0
 
-        # Apply 2D DCT in 8x8 blocks
+        # Image should already be 8x8 aligned
         h, w = y_channel.shape
+        assert h % 8 == 0 and w % 8 == 0, f"Image must be 8x8 aligned, got {h}x{w}"
 
         # Compute DCT for each 8x8 block
-        dct_coeffs = np.zeros_like(y_channel)
+        dct_coeffs = np.zeros((h, w), dtype=np.float32)
         for i in range(0, h, 8):
             for j in range(0, w, 8):
                 block = y_channel[i:i+8, j:j+8]
@@ -289,8 +290,14 @@ class DTDPredictor:
         full_mask = full_mask / np.maximum(count_map, 1)
         final_mask = (full_mask > 0.5).astype(np.uint8)
 
-        # Crop back to TRUE original size (before any padding)
-        final_mask = final_mask[:true_orig_h, :true_orig_w]
+        # Pad mask back to true original size if it was 8x8 aligned smaller
+        if final_mask.shape[0] < true_orig_h or final_mask.shape[1] < true_orig_w:
+            padded_mask = np.zeros((true_orig_h, true_orig_w), dtype=np.uint8)
+            padded_mask[:final_mask.shape[0], :final_mask.shape[1]] = final_mask
+            final_mask = padded_mask
+        else:
+            # Crop if somehow larger (shouldn't happen)
+            final_mask = final_mask[:true_orig_h, :true_orig_w]
 
         # Create heatmap overlay with original image (no padding)
         heatmap = self.create_heatmap(im_orig_np, final_mask)