get_affine_transform is incorrect when scale_w!=scale_h

[imistyrain]

in transforms.py

    scale_tmp = scale * 200.0
    src_w = scale_tmp[0]
    dst_w = output_size[0]
    dst_h = output_size[1]

    rot_rad = np.pi * rot / 180
    src_dir = get_dir([0, src_w * -0.5], rot_rad)
    dst_dir = np.array([0, dst_w * -0.5], np.float32)

    src = np.zeros((3, 2), dtype=np.float32)
    dst = np.zeros((3, 2), dtype=np.float32)
    src[0, :] = center + scale_tmp * shift
    src[1, :] = center + src_dir + scale_tmp * shift
    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir

src_w is used as src_h.
and the correct one should be

    scale_tmp = scale * 200.0
    if scale[1] > scale[0]:
        src_h = scale_tmp[1]
    else:
        src_w = scale_tmp[0]
    dst_w = output_size[0]
    dst_h = output_size[1]
  
    rot_rad = np.pi * rot / 180
    if scale[1] > scale[0]:
        src_dir = get_dir([0, src_h * -0.5], rot_rad)
        dst_dir = np.array([0, dst_h * -0.5], np.float32)
    else:
        src_dir = get_dir([0, src_w * -0.5], rot_rad)
        dst_dir = np.array([0, dst_w * -0.5], np.float32)
  
    src = np.zeros((3, 2), dtype=np.float32)
    dst = np.zeros((3, 2), dtype=np.float32)
    src[0, :] = center + scale_tmp * shift
    src[1, :] = center + src_dir + scale_tmp * shift
    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
  
    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])

[Jackqu]

I think you are right！