Add FITS Muxer

Signed-off-by: Paras Chadha <paraschadha18@gmail.com>

4 files changed, 187 insertions, 1 deletions

diff --git a/libavformat/Makefile b/libavformat/Makefile
index dfae160bee..36f5839aa8 100644
--- a/libavformat/Makefile
+++ b/libavformat/Makefile
@@ -165,6 +165,7 @@ OBJS-$(CONFIG_FIFO_MUXER)                += fifo.o
 OBJS-$(CONFIG_FILMSTRIP_DEMUXER)         += filmstripdec.o
 OBJS-$(CONFIG_FILMSTRIP_MUXER)           += filmstripenc.o
 OBJS-$(CONFIG_FITS_DEMUXER)              += fitsdec.o
+OBJS-$(CONFIG_FITS_MUXER)                += fitsenc.o
 OBJS-$(CONFIG_FLAC_DEMUXER)              += flacdec.o rawdec.o \
                                             flac_picture.o   \
                                             oggparsevorbis.o \
diff --git a/libavformat/allformats.c b/libavformat/allformats.c
index f23a9a06e5..cb09a60e6f 100644
--- a/libavformat/allformats.c
+++ b/libavformat/allformats.c
@@ -121,7 +121,7 @@ static void register_all(void)
     REGISTER_MUXDEMUX(FFMETADATA,       ffmetadata);
     REGISTER_MUXER   (FIFO,             fifo);
     REGISTER_MUXDEMUX(FILMSTRIP,        filmstrip);
-    REGISTER_DEMUXER (FITS,             fits);
+    REGISTER_MUXDEMUX(FITS,             fits);
     REGISTER_MUXDEMUX(FLAC,             flac);
     REGISTER_DEMUXER (FLIC,             flic);
     REGISTER_MUXDEMUX(FLV,              flv);
diff --git a/libavformat/fitsenc.c b/libavformat/fitsenc.c
new file mode 100644
index 0000000000..0dcdcdfb04
--- /dev/null
+++ b/libavformat/fitsenc.c
@@ -0,0 +1,183 @@
+/*
+ * FITS muxer
+ * Copyright (c) 2017 Paras Chadha
+ *
+ * This file is part of FFmpeg.
+ *
+ * FFmpeg is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * FFmpeg 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
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with FFmpeg; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+/**
+ * @file
+ * FITS muxer.
+ */
+
+#include "internal.h"
+
+typedef struct FITSContext {
+    int first_image;
+} FITSContext;
+
+static int fits_write_header(AVFormatContext *s)
+{
+    FITSContext *fitsctx = s->priv_data;
+    fitsctx->first_image = 1;
+    return 0;
+}
+
+/**
+ * Write one header line comprising of keyword and value(int)
+ * @param s AVFormat Context
+ * @param keyword pointer to the char array in which keyword is stored
+ * @param value the value corresponding to the keyword
+ * @param lines_written to keep track of lines written so far
+ * @return 0
+ */
+static int write_keyword_value(AVFormatContext *s, const char *keyword, int value, int *lines_written)
+{
+    int len, ret;
+    uint8_t header[80];
+
+    len = strlen(keyword);
+    memset(header, ' ', sizeof(header));
+    memcpy(header, keyword, len);
+
+    header[8] = '=';
+    header[9] = ' ';
+
+    ret = snprintf(header + 10, 70, "%d", value);
+    header[ret + 10] = ' ';
+
+    avio_write(s->pb, header, sizeof(header));
+    *lines_written += 1;
+    return 0;
+}
+
+static int write_image_header(AVFormatContext *s)
+{
+    AVStream *st = s->streams[0];
+    AVCodecParameters *encctx = st->codecpar;
+    FITSContext *fitsctx = s->priv_data;
+    uint8_t buffer[80];
+    int bitpix, naxis, naxis3 = 1, bzero = 0, rgb = 0, lines_written = 0, lines_left;
+
+    switch (encctx->format) {
+        case AV_PIX_FMT_GRAY8:
+            bitpix = 8;
+            naxis = 2;
+            break;
+        case AV_PIX_FMT_GRAY16BE:
+            bitpix = 16;
+            naxis = 2;
+            bzero = 32768;
+            break;
+        case AV_PIX_FMT_GBRP:
+        case AV_PIX_FMT_GBRAP:
+            bitpix = 8;
+            naxis = 3;
+            rgb = 1;
+            if (encctx->format == AV_PIX_FMT_GBRP) {
+                naxis3 = 3;
+            } else {
+                naxis3 = 4;
+            }
+            break;
+        case AV_PIX_FMT_GBRP16BE:
+        case AV_PIX_FMT_GBRAP16BE:
+            bitpix = 16;
+            naxis = 3;
+            rgb = 1;
+            if (encctx->format == AV_PIX_FMT_GBRP16BE) {
+                naxis3 = 3;
+            } else {
+                naxis3 = 4;
+            }
+            bzero = 32768;
+            break;
+    }
+
+    if (fitsctx->first_image) {
+        memcpy(buffer, "SIMPLE  = ", 10);
+        memset(buffer + 10, ' ', 70);
+        buffer[29] = 'T';
+        avio_write(s->pb, buffer, sizeof(buffer));
+    } else {
+        memcpy(buffer, "XTENSION= 'IMAGE   '", 20);
+        memset(buffer + 20, ' ', 60);
+        avio_write(s->pb, buffer, sizeof(buffer));
+    }
+    lines_written++;
+
+    write_keyword_value(s, "BITPIX", bitpix, &lines_written);         // no of bits per pixel
+    write_keyword_value(s, "NAXIS", naxis, &lines_written);           // no of dimensions of image
+    write_keyword_value(s, "NAXIS1", encctx->width, &lines_written);   // first dimension i.e. width
+    write_keyword_value(s, "NAXIS2", encctx->height, &lines_written);  // second dimension i.e. height
+
+    if (rgb)
+        write_keyword_value(s, "NAXIS3", naxis3, &lines_written);     // third dimension to store RGBA planes
+
+    if (!fitsctx->first_image) {
+        write_keyword_value(s, "PCOUNT", 0, &lines_written);
+        write_keyword_value(s, "GCOUNT", 1, &lines_written);
+    } else {
+        fitsctx->first_image = 0;
+    }
+
+    /*
+     * Since FITS does not support unsigned 16 bit integers,
+     * BZERO = 32768 is used to store unsigned 16 bit integers as
+     * signed integers so that it can be read properly.
+     */
+    if (bitpix == 16)
+        write_keyword_value(s, "BZERO", bzero, &lines_written);
+
+    if (rgb) {
+        memcpy(buffer, "CTYPE3  = 'RGB     '", 20);
+        memset(buffer + 20, ' ', 60);
+        avio_write(s->pb, buffer, sizeof(buffer));
+        lines_written++;
+    }
+
+    memcpy(buffer, "END", 3);
+    memset(buffer + 3, ' ', 77);
+    avio_write(s->pb, buffer, sizeof(buffer));
+    lines_written++;
+
+    lines_left = ((lines_written + 35) / 36) * 36 - lines_written;
+    memset(buffer, ' ', 80);
+    while (lines_left > 0) {
+        avio_write(s->pb, buffer, sizeof(buffer));
+        lines_left--;
+    }
+    return 0;
+}
+
+static int fits_write_packet(AVFormatContext *s, AVPacket *pkt)
+{
+    write_image_header(s);
+    avio_write(s->pb, pkt->data, pkt->size);
+    return 0;
+}
+
+AVOutputFormat ff_fits_muxer = {
+    .name         = "fits",
+    .long_name    = NULL_IF_CONFIG_SMALL("Flexible Image Transport System"),
+    .extensions   = "fits",
+    .priv_data_size = sizeof(FITSContext),
+    .audio_codec  = AV_CODEC_ID_NONE,
+    .video_codec  = AV_CODEC_ID_FITS,
+    .write_header = fits_write_header,
+    .write_packet = fits_write_packet,
+};
diff --git a/libavformat/img2enc.c b/libavformat/img2enc.c
index 1297b1aaba..87b5ec2317 100644
--- a/libavformat/img2enc.c
+++ b/libavformat/img2enc.c
@@ -62,6 +62,8 @@ static int write_header(AVFormatContext *s)
 
     if (st->codecpar->codec_id == AV_CODEC_ID_GIF) {
         img->muxer = "gif";
+    } else if (st->codecpar->codec_id == AV_CODEC_ID_FITS) {
+        img->muxer = "fits";
     } else if (st->codecpar->codec_id == AV_CODEC_ID_RAWVIDEO) {
         const char *str = strrchr(img->path, '.');
         img->split_planes =     str