Skip to content

Commit

Permalink
feat: load all HEIC/AVIF aux channels (bilinear resize to main img)
Browse files Browse the repository at this point in the history
  • Loading branch information
Tom94 committed Feb 25, 2025
1 parent 9eb9fce commit 081c103
Show file tree
Hide file tree
Showing 4 changed files with 90 additions and 37 deletions.
2 changes: 1 addition & 1 deletion include/tev/imageio/ImageLoader.h
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ class ImageLoader {
static const std::vector<std::unique_ptr<ImageLoader>>& getLoaders();

protected:
static std::vector<Channel> makeNChannels(int numChannels, const nanogui::Vector2i& size);
static std::vector<Channel> makeNChannels(int numChannels, const nanogui::Vector2i& size, const std::string& namePrefix = "");
};

} // namespace tev
27 changes: 2 additions & 25 deletions src/imageio/GainMap.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -27,28 +27,7 @@ namespace tev {

Task<void> applyAppleGainMap(ImageData& image, const ImageData& gainMap, int priority, const AppleMakerNote& amn) {
auto size = image.channels[0].size();
auto gainMapSize = gainMap.channels[0].size();

// Assumption: gain maps are always a power of 2 smaller than the image, give or take rounding for odd image dimensions.
// Algorithm: find first power of 2 downscale of the image that is smaller than the gain map.
uint32_t shift = 0;
while ((size.x() >> shift) > gainMapSize.x() && (size.y() >> shift) > gainMapSize.y()) {
++shift;
}

// Warn if final gainmapsize obtained by shifting is off my more than 1 pixels and don't use the gainmap.
if (abs((int)(size.x() >> shift) - (int)gainMapSize.x()) > 1 || abs((int)(size.y() >> shift) - (int)gainMapSize.y()) > 1) {
tlog::warning() << fmt::format(
"Gain map size {}x{} does not match {}-downscaled image size {}x{}",
gainMapSize.x(),
gainMapSize.y(),
shift,
size.x() >> shift,
size.y() >> shift
);

co_return;
}
TEV_ASSERT(size == gainMap.channels[0].size(), "Image and gain map must have the same size");

// Apply gain map per https://developer.apple.com/documentation/appkit/applying-apple-hdr-effect-to-your-photos
float headroom = 1.0f;
Expand Down Expand Up @@ -86,10 +65,8 @@ Task<void> applyAppleGainMap(ImageData& image, const ImageData& gainMap, int pri
[&](int y) {
for (int x = 0; x < size.x(); ++x) {
size_t i = y * (size_t)size.x() + x;
size_t gmi = (x >> shift) + (y >> shift) * gainMapSize.x();

for (int c = 0; c < 3; ++c) {
image.channels[c].at(i) *= (1.0f + (headroom - 1.0f) * gainMap.channels[0].at(gmi));
image.channels[c].at(i) *= (1.0f + (headroom - 1.0f) * gainMap.channels[0].at(i));
}
}
},
Expand Down
94 changes: 85 additions & 9 deletions src/imageio/HeifImageLoader.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

#include <tev/Common.h>
#include <tev/ThreadPool.h>
#include <tev/imageio/AppleMakerNote.h>
#include <tev/imageio/GainMap.h>
Expand Down Expand Up @@ -76,7 +77,7 @@ bool HeifImageLoader::canLoadFile(istream& iStream) const {
return heif_check_filetype(header, 12) == heif_filetype_yes_supported;
}

Task<vector<ImageData>> HeifImageLoader::load(istream& iStream, const fs::path&, const string&, int priority) const {
Task<vector<ImageData>> HeifImageLoader::load(istream& iStream, const fs::path&, const string& channelSelector, int priority) const {
vector<ImageData> result;

iStream.seekg(0, ios_base::end);
Expand Down Expand Up @@ -135,7 +136,7 @@ Task<vector<ImageData>> HeifImageLoader::load(istream& iStream, const fs::path&,

ScopeGuard handleGuard{[handle] { heif_image_handle_release(handle); }};

auto decodeImage = [priority](heif_image_handle* imgHandle) -> Task<ImageData> {
auto decodeImage = [priority](heif_image_handle* imgHandle, const Vector2i& targetSize = {0}, const string& namePrefix = "") -> Task<ImageData> {
ImageData resultData;

int numChannels = heif_image_handle_has_alpha_channel(imgHandle) ? 4 : 3;
Expand All @@ -158,6 +159,13 @@ Task<vector<ImageData>> HeifImageLoader::load(istream& iStream, const fs::path&,

ScopeGuard imgGuard{[img] { heif_image_release(img); }};

// if (targetSize.x() != 0 && size != targetSize) {
// heif_image* scaledImg;
// heif_image_scale_image(img, &scaledImg, targetSize.x(), targetSize.y(), nullptr);
// heif_image_release(img);
// img = scaledImg;
// }

const int bitsPerPixel = heif_image_get_bits_per_pixel_range(img, heif_channel_interleaved);
const float channelScale = 1.0f / float((1 << bitsPerPixel) - 1);

Expand All @@ -167,7 +175,7 @@ Task<vector<ImageData>> HeifImageLoader::load(istream& iStream, const fs::path&,
throw invalid_argument{"Faild to get image data."};
}

resultData.channels = makeNChannels(numChannels, size);
resultData.channels = makeNChannels(numChannels, size, namePrefix);

auto getCmsTransform = [&imgHandle, &numChannels, &resultData]() {
size_t profileSize = heif_image_handle_get_raw_color_profile_size(imgHandle);
Expand Down Expand Up @@ -322,6 +330,7 @@ Task<vector<ImageData>> HeifImageLoader::load(istream& iStream, const fs::path&,

// Read main image
result.emplace_back(co_await decodeImage(handle));
ImageData& mainImage = result.front();

auto findAppleMakerNote = [&]() -> unique_ptr<AppleMakerNote> {
// Extract EXIF metadata
Expand Down Expand Up @@ -372,6 +381,67 @@ Task<vector<ImageData>> HeifImageLoader::load(istream& iStream, const fs::path&,

auto amn = findAppleMakerNote();

auto resizeImage = [priority](ImageData& resultData, const Vector2i& targetSize, const string& namePrefix) -> Task<void> {
Vector2i size = resultData.channels.front().size();
if (size == targetSize) {
co_return;
}

int numChannels = (int)resultData.channels.size();

ImageData scaledResultData;
scaledResultData.hasPremultipliedAlpha = resultData.hasPremultipliedAlpha;
scaledResultData.channels = makeNChannels(numChannels, targetSize, namePrefix);

auto& srcChannels = resultData.channels;
auto& dstChannels = scaledResultData.channels;

co_await ThreadPool::global().parallelForAsync<int>(
0,
targetSize.y(),
[&](int dstY) {
const float scaleX = (float)size.x() / targetSize.x();
const float scaleY = (float)size.y() / targetSize.y();

for (int dstX = 0; dstX < targetSize.x(); ++dstX) {
float srcX = (dstX + 0.5f) * scaleX - 0.5f;
float srcY = (dstY + 0.5f) * scaleY - 0.5f;

int x0 = std::max((int)std::floor(srcX), 0);
int y0 = std::max((int)std::floor(srcY), 0);
int x1 = std::min(x0 + 1, size.x() - 1);
int y1 = std::min(y0 + 1, size.y() - 1);

float wx1 = srcX - x0;
float wy1 = srcY - y0;
float wx0 = 1.0f - wx1;
float wy0 = 1.0f - wy1;

size_t dstIdx = dstY * (size_t)targetSize.x() + dstX;

size_t srcIdx00 = y0 * (size_t)size.x() + x0;
size_t srcIdx01 = y0 * (size_t)size.x() + x1;
size_t srcIdx10 = y1 * (size_t)size.x() + x0;
size_t srcIdx11 = y1 * (size_t)size.x() + x1;

for (int c = 0; c < numChannels; ++c) {
float p00 = srcChannels[c].at(srcIdx00);
float p01 = srcChannels[c].at(srcIdx01);
float p10 = srcChannels[c].at(srcIdx10);
float p11 = srcChannels[c].at(srcIdx11);

float interpolated = wy0 * (wx0 * p00 + wx1 * p01) + wy1 * (wx0 * p10 + wx1 * p11);
dstChannels[c].at(dstIdx) = interpolated;
}
}
},
priority
);

resultData = std::move(scaledResultData);
co_return;
};

// Read auxiliary images
int num_aux = heif_image_handle_get_number_of_auxiliary_images(handle, 0);
if (num_aux > 0) {
Expand All @@ -394,17 +464,23 @@ Task<vector<ImageData>> HeifImageLoader::load(istream& iStream, const fs::path&,
}

ScopeGuard typeGuard{[auxImgHandle, &auxType] { heif_image_handle_release_auxiliary_type(auxImgHandle, &auxType); }};
string auxTypeStr = auxType ? auxType : "unknown";
string auxLayerName = auxType ? fmt::format("{}.", auxType) : fmt::format("{}.", num_aux);
replace(auxLayerName.begin(), auxLayerName.end(), ':', '.');

if (!matchesFuzzy(auxLayerName, channelSelector)) {
continue;
}

// TODO: Better handling of auxiliary images than to just decode them and list them as separate images
// result.emplace_back(co_await decodeImage(auxImgHandle));
auto auxImgData = co_await decodeImage(auxImgHandle, mainImage.channels.front().size(), auxLayerName);
co_await resizeImage(auxImgData, mainImage.channels.front().size(), auxLayerName);
mainImage.channels.insert(mainImage.channels.end(), auxImgData.channels.begin(), auxImgData.channels.end());

// If we found an apple-style gainmap, apply it to the main image.
if (amn && auxTypeStr.find("apple") != string::npos && auxTypeStr.find("hdrgainmap") != string::npos) {
tlog::debug() << fmt::format("Found hdrgainmap: {}", auxTypeStr);
if (amn && auxLayerName.find("apple") != string::npos && auxLayerName.find("hdrgainmap") != string::npos) {
tlog::debug() << fmt::format("Found hdrgainmap: {}", auxLayerName);
co_await applyAppleGainMap(
result.front(), // primary image
co_await decodeImage(auxImgHandle),
auxImgData,
priority,
*amn
);
Expand Down
4 changes: 2 additions & 2 deletions src/imageio/ImageLoader.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ const vector<unique_ptr<ImageLoader>>& ImageLoader::getLoaders() {
return imageLoaders;
}

vector<Channel> ImageLoader::makeNChannels(int numChannels, const Vector2i& size) {
vector<Channel> ImageLoader::makeNChannels(int numChannels, const Vector2i& size, const string& namePrefix) {
vector<Channel> channels;

size_t numPixels = (size_t)size.x() * size.y();
Expand All @@ -73,7 +73,7 @@ vector<Channel> ImageLoader::makeNChannels(int numChannels, const Vector2i& size
if (numChannels > 1) {
const vector<string> channelNames = {"R", "G", "B", "A"};
for (int c = 0; c < numChannels; ++c) {
string name = c < (int)channelNames.size() ? channelNames[c] : to_string(c);
string name = namePrefix + (c < (int)channelNames.size() ? channelNames[c] : to_string(c));

// We assume that the channels are interleaved.
channels.emplace_back(name, size, data, c, 4);
Expand Down

0 comments on commit 081c103

Please sign in to comment.