From f4b128df14255ba7f6a834054c379e59b8d1d659 Mon Sep 17 00:00:00 2001
From: guoyongzhi <guoyongzhi@megvii.com>
Date: Tue, 25 Feb 2025 13:32:26 +0800
Subject: [PATCH] fix(qtrees): Update PaddedMat constructor to handle cases
 where kernel size is smaller than the input image

---
 Project.toml  | 2 +-
 src/qtrees.jl | 9 ++++++---
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/Project.toml b/Project.toml
index bce7664..8616f40 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Stuffing"
 uuid = "4175e07e-e5b7-423e-8796-3ea7f6d48281"
 authors = ["guoyongzhi <guo-yong-zhi@outlook.com>"]
-version = "0.10.2"
+version = "0.10.3"
 
 [deps]
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
diff --git a/src/qtrees.jl b/src/qtrees.jl
index 4ad8687..e7958f4 100644
--- a/src/qtrees.jl
+++ b/src/qtrees.jl
@@ -115,8 +115,12 @@ mutable struct PaddedMat{T <: AbstractMatrix{UInt8}} <: AbstractMatrix{UInt8}
 end
 
 function PaddedMat(l::T, sz::Tuple{Int,Int}=size(l), rshift=0, cshift=0; default=0x00) where {T <: AbstractMatrix{UInt8}}
-    m = PaddedMat{T}(size(l), sz, rshift, cshift; default=default)
-    m.kernel[2:end - 2, 2:end - 2] .= l
+    (sz[1] >= size(l, 1) && sz[2] >= size(l, 2)) || (@warn "PaddedMat cut off: $(size(l))>$sz")
+    ksz = min.(size(l), sz)
+    b1, b2 = (size(l) .- ksz) .÷ 2
+    e1, e2 = (b1, b2) .+ ksz
+    m = PaddedMat{T}(ksz, sz, rshift, cshift; default=default)
+    m.kernel[2:end - 2, 2:end - 2] .= l[b1+1:e1, b2+1:e2]
     m
 end
 function PaddedMat{T}(kernelsz::Tuple{Int,Int}, sz::Tuple{Int,Int}, 
@@ -193,7 +197,6 @@ function ShiftedQTree(pic::PaddedMat{T}) where T
 end
 function ShiftedQTree(pic::AbstractMatrix{UInt8}, sz::Integer; default=EMPTY)
     @assert isinteger(log2(sz))
-    (sz >= size(pic, 1) && sz >= size(pic, 2)) || (@warn "ShiftedQTree cut off: $(size(pic))>$sz")
     ShiftedQTree(PaddedMat(pic, (sz, sz), default=default))
 end
 function ShiftedQTree(pic::AbstractMatrix{UInt8}; default=EMPTY)