omdb/image.cpp

/*
 * Copyright (c) 2017-2023, Stefan Haustein, Aaron Liu
 *
 *     This file is free software: you may copy, redistribute and/or modify it
 *     under the terms of the GNU General Public License as published by the
 *     Free Software Foundation, either version 3 of the License, or (at your
 *     option) any later version.
 *
 *     This file is distributed in the hope that it will be useful, but
 *     WITHOUT ANY WARRANTY; without even the implied warranty of
 *     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 *     General Public License for more details.
 *
 *     You should have received a copy of the GNU General Public License
 *     along with this program.  If not, see <https://www.gnu.org/licenses/>.
 *
 * Alternatively, you may copy, redistribute and/or modify this file under
 * the terms of the Apache License, version 2.0:
 *
 *     Licensed under the Apache License, Version 2.0 (the "License");
 *     you may not use this file except in compliance with the License.
 *     You may obtain a copy of the License at
 *
 *         https://www.apache.org/licenses/LICENSE-2.0
 *
 *     Unless required by applicable law or agreed to in writing, software
 *     distributed under the License is distributed on an "AS IS" BASIS,
 *     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *     See the License for the specific language governing permissions and
 *     limitations under the License.
 */

#include <algorithm>
#include <array>
#include <cmath>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>

#include "tiv_lib.h"

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#define STB_IMAGE_RESIZE_IMPLEMENTATION
#include "stb_image_resize2.h"

#include <cstring>
#include <sys/ioctl.h>
#include <unistd.h>

// Program exit code constants compatible with sysexits.h.
#define EXITCODE_OK 0
#define EXITCODE_COMMAND_LINE_USAGE_ERROR 64
#define EXITCODE_DATA_FORMAT_ERROR 65
#define EXITCODE_NO_INPUT_ERROR 66

void printTermColor(const int &flags, int r, int g, int b) {
  r = clamp_byte(r);
  g = clamp_byte(g);
  b = clamp_byte(b);

  bool bg = (flags & FLAG_BG) != 0;

  if ((flags & FLAG_MODE_256) == 0) {
    std::cout << (bg ? "\x1b[48;2;" : "\x1b[38;2;") << r << ';' << g << ';' << b
              << 'm';
    return;
  }

  int ri = best_index(r, COLOR_STEPS, COLOR_STEP_COUNT);
  int gi = best_index(g, COLOR_STEPS, COLOR_STEP_COUNT);
  int bi = best_index(b, COLOR_STEPS, COLOR_STEP_COUNT);

  int rq = COLOR_STEPS[ri];
  int gq = COLOR_STEPS[gi];
  int bq = COLOR_STEPS[bi];

  int gray =
      static_cast<int>(std::round(r * 0.2989f + g * 0.5870f + b * 0.1140f));

  int gri = best_index(gray, GRAYSCALE_STEPS, GRAYSCALE_STEP_COUNT);
  int grq = GRAYSCALE_STEPS[gri];

  int color_index;
  if (0.3 * sqr(rq - r) + 0.59 * sqr(gq - g) + 0.11 * sqr(bq - b) <
      0.3 * sqr(grq - r) + 0.59 * sqr(grq - g) + 0.11 * sqr(grq - b)) {
    color_index = 16 + 36 * ri + 6 * gi + bi;
  } else {
    color_index = 232 + gri; // 1..24 -> 232..255
  }
  std::cout << (bg ? "\x1B[48;5;" : "\u001B[38;5;") << color_index << "m";
}

void printCodepoint(int codepoint) {
  if (codepoint < 128) {
    std::cout << static_cast<char>(codepoint);
  } else if (codepoint < 0x7ff) {
    std::cout << static_cast<char>(0xc0 | (codepoint >> 6));
    std::cout << static_cast<char>(0x80 | (codepoint & 0x3f));
  } else if (codepoint < 0xffff) {
    std::cout << static_cast<char>(0xe0 | (codepoint >> 12));
    std::cout << static_cast<char>(0x80 | ((codepoint >> 6) & 0x3f));
    std::cout << static_cast<char>(0x80 | (codepoint & 0x3f));
  } else if (codepoint < 0x10ffff) {
    std::cout << static_cast<char>(0xf0 | (codepoint >> 18));
    std::cout << static_cast<char>(0x80 | ((codepoint >> 12) & 0x3f));
    std::cout << static_cast<char>(0x80 | ((codepoint >> 6) & 0x3f));
    std::cout << static_cast<char>(0x80 | (codepoint & 0x3f));
  } else {
    std::cerr << "ERROR";
  }
}

void printImage(const unsigned char *data, int width, int height, int channels,
                const int &flags) {
  GetPixelFunction get_pixel = [&](int x, int y) -> unsigned long {
    int index = (y * width + x) * channels;
    return (((unsigned long)data[index]) << 16) |
           (((unsigned long)data[index + 1]) << 8) |
           (((unsigned long)data[index + 2]));
  };

  CharData lastCharData;
  for (int y = 0; y <= height - 8; y += 8) {
    for (int x = 0; x <= width - 4; x += 4) {
      CharData charData =
          flags & FLAG_NOOPT
              ? createCharData(get_pixel, x, y, 0x2584, 0x0000ffff)
              : findCharData(get_pixel, x, y, flags);
      if (x == 0 || charData.bgColor != lastCharData.bgColor)
        printTermColor(flags | FLAG_BG, charData.bgColor[0],
                       charData.bgColor[1], charData.bgColor[2]);
      if (x == 0 || charData.fgColor != lastCharData.fgColor)
        printTermColor(flags | FLAG_FG, charData.fgColor[0],
                       charData.fgColor[1], charData.fgColor[2]);
      printCodepoint(charData.codePoint);
      lastCharData = charData;
    }
    std::cout << "\x1b[0m" << std::endl;
  }
}

struct size {
  size(unsigned int in_width, unsigned int in_height)
      : width(in_width), height(in_height) {}
  unsigned int width;
  unsigned int height;
  size scaled(double scale) { return size(width * scale, height * scale); }
  size fitted_within(size container) {
    double scale = std::min(container.width / static_cast<double>(width),
                            container.height / static_cast<double>(height));
    return scaled(scale);
  }
};
std::ostream &operator<<(std::ostream &stream, size sz) {
  stream << sz.width << "x" << sz.height;
  return stream;
}

static void render_and_print_image(unsigned char *data, int width, int height,
                                   int channels) {
  int max_width = 80;
  int max_height = 24;
  struct winsize w;
  char err_buf[LINE_MAX];

  if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &w) == -1) {
    strerror_r(errno, err_buf, sizeof err_buf);
    std::cerr << "Couldn't get terminal dimensions: " << err_buf << std::endl;
  }

  // TODO: Maybe make this configuratble, but quality suffers at smaller sizes.
  max_width = w.ws_col * 4;
  max_height = w.ws_row * 8;

  unsigned char *output_data = data;
  int output_width = width;
  int output_height = height;

  std::vector<unsigned char> resized_data;

  if (width > max_width || height > max_height) {
    size new_size =
        size(width, height).fitted_within(size(max_width, max_height));
    output_width = new_size.width;
    output_height = new_size.height;

    resized_data.resize(output_width * output_height * channels);
    stbir_resize_uint8_linear(data, width, height, 0, resized_data.data(),
                              output_width, output_height, 0,
                              (stbir_pixel_layout)channels);
    output_data = resized_data.data();
  }

  int flags = 0;
  printImage(output_data, output_width, output_height, channels, flags);
}

extern "C" {
void process_image(const char *path) {
  int width, height, channels;
  unsigned char *data = stbi_load(path, &width, &height, &channels, 3);
  if (!data) {
    std::cerr << "Error: '" << path << "' could not be loaded" << std::endl;
    return;
  }
  // Force 3 channels for now
  channels = 3;

  render_and_print_image(data, width, height, channels);

  stbi_image_free(data);
}

void process_image_mem(const unsigned char *buffer, size_t len) {
  int width, height, channels;
  unsigned char *data =
      stbi_load_from_memory(buffer, (int)len, &width, &height, &channels, 3);
  if (!data) {
    std::cerr << "Error: Could not load image from memory" << std::endl;
    return;
  }
  // Force 3 channels for now
  channels = 3;

  render_and_print_image(data, width, height, channels);

  stbi_image_free(data);
}
}