49a40c461acfdd151ce83ebebac22f0ff9d24a24
[blender.git] / source / blender / imbuf / intern / indexer.c
1 /*
2  * ***** BEGIN GPL LICENSE BLOCK *****
3  *
4  * This program is free software; you can redistribute it and/or
5  * modify it under the terms of the GNU General Public License
6  * as published by the Free Software Foundation; either version 2
7  * of the License, or (at your option) any later version.
8  *
9  * This program is distributed in the hope that it will be useful,
10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12  * GNU General Public License for more details.
13  *
14  * You should have received a copy of the GNU General Public License
15  * along with this program; if not, write to the Free Software Foundation,
16  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
17  *
18  * Peter Schlaile <peter [at] schlaile [dot] de> 2011
19  *
20  * Contributor(s): none yet.
21  *
22  * ***** END GPL LICENSE BLOCK *****
23  */
24
25 /** \file blender/imbuf/intern/indexer.c
26  *  \ingroup imbuf
27  */
28
29 #include <stdlib.h>
30
31 #include "MEM_guardedalloc.h"
32
33 #include "BLI_utildefines.h"
34 #include "BLI_endian_switch.h"
35 #include "BLI_path_util.h"
36 #include "BLI_string.h"
37 #include "BLI_path_util.h"
38 #include "BLI_fileops.h"
39 #include "BLI_math_base.h"
40
41 #include "IMB_indexer.h"
42 #include "IMB_anim.h"
43 #include "imbuf.h"
44
45 #include "MEM_guardedalloc.h"
46 #include "DNA_userdef_types.h"
47 #include "BKE_global.h"
48
49 #ifdef WITH_AVI
50 #  include "AVI_avi.h"
51 #endif
52
53 #ifdef WITH_FFMPEG
54 #  include "ffmpeg_compat.h"
55 #endif
56
57
58 static const char magic[] = "BlenMIdx";
59 static const char temp_ext[] = "_part";
60
61 static const int proxy_sizes[] = { IMB_PROXY_25, IMB_PROXY_50, IMB_PROXY_75,
62                                  IMB_PROXY_100 };
63 static const float proxy_fac[] = { 0.25, 0.50, 0.75, 1.00 };
64
65 #ifdef WITH_FFMPEG
66 static int tc_types[] = {IMB_TC_RECORD_RUN,
67                          IMB_TC_FREE_RUN,
68                          IMB_TC_INTERPOLATED_REC_DATE_FREE_RUN,
69                          IMB_TC_RECORD_RUN_NO_GAPS};
70 #endif
71
72 #define INDEX_FILE_VERSION 1
73
74 /* ----------------------------------------------------------------------
75  * - time code index functions
76  * ---------------------------------------------------------------------- */
77
78 anim_index_builder *IMB_index_builder_create(const char *name)
79 {
80
81         anim_index_builder *rv = MEM_callocN(sizeof(struct anim_index_builder),
82                                              "index builder");
83
84         fprintf(stderr, "Starting work on index: %s\n", name);
85
86         BLI_strncpy(rv->name, name, sizeof(rv->name));
87         BLI_strncpy(rv->temp_name, name, sizeof(rv->temp_name));
88
89         strcat(rv->temp_name, temp_ext);
90
91         BLI_make_existing_file(rv->temp_name);
92
93         rv->fp = BLI_fopen(rv->temp_name, "wb");
94
95         if (!rv->fp) {
96                 fprintf(stderr, "Couldn't open index target: %s! "
97                         "Index build broken!\n", rv->temp_name);
98                 MEM_freeN(rv);
99                 return NULL;
100         }
101
102         fprintf(rv->fp, "%s%c%.3d", magic, (ENDIAN_ORDER == B_ENDIAN) ? 'V' : 'v',
103                 INDEX_FILE_VERSION);
104
105         return rv;
106 }
107
108 void IMB_index_builder_add_entry(anim_index_builder *fp,
109                                  int frameno, unsigned long long seek_pos,
110                                  unsigned long long seek_pos_dts,
111                                  unsigned long long pts)
112 {
113         fwrite(&frameno, sizeof(int), 1, fp->fp);
114         fwrite(&seek_pos, sizeof(unsigned long long), 1, fp->fp);
115         fwrite(&seek_pos_dts, sizeof(unsigned long long), 1, fp->fp);
116         fwrite(&pts, sizeof(unsigned long long), 1, fp->fp);
117 }
118
119 void IMB_index_builder_proc_frame(anim_index_builder *fp,
120                                   unsigned char *buffer,
121                                   int data_size,
122                                   int frameno, unsigned long long seek_pos,
123                                   unsigned long long seek_pos_dts,
124                                   unsigned long long pts)
125 {
126         if (fp->proc_frame) {
127                 anim_index_entry e;
128                 e.frameno = frameno;
129                 e.seek_pos = seek_pos;
130                 e.seek_pos_dts = seek_pos_dts;
131                 e.pts = pts;
132
133                 fp->proc_frame(fp, buffer, data_size, &e);
134         }
135         else {
136                 IMB_index_builder_add_entry(fp, frameno, seek_pos,
137                                             seek_pos_dts, pts);
138         }
139 }
140
141 void IMB_index_builder_finish(anim_index_builder *fp, int rollback)
142 {
143         if (fp->delete_priv_data) {
144                 fp->delete_priv_data(fp);
145         }
146
147         fclose(fp->fp);
148         
149         if (rollback) {
150                 unlink(fp->temp_name);
151         }
152         else {
153                 unlink(fp->name);
154                 BLI_rename(fp->temp_name, fp->name);
155         }
156
157         MEM_freeN(fp);
158 }
159
160 struct anim_index *IMB_indexer_open(const char *name)
161 {
162         char header[13];
163         struct anim_index *idx;
164         FILE *fp = BLI_fopen(name, "rb");
165         int i;
166
167         if (!fp) {
168                 return NULL;
169         }
170
171         if (fread(header, 12, 1, fp) != 1) {
172                 fclose(fp);
173                 return NULL;
174         }
175
176         header[12] = 0;
177
178         if (memcmp(header, magic, 8) != 0) {
179                 fclose(fp);
180                 return NULL;
181         }
182
183         if (atoi(header + 9) != INDEX_FILE_VERSION) {
184                 fclose(fp);
185                 return NULL;
186         }
187
188         idx = MEM_callocN(sizeof(struct anim_index), "anim_index");
189
190         BLI_strncpy(idx->name, name, sizeof(idx->name));
191         
192         fseek(fp, 0, SEEK_END);
193
194         idx->num_entries = (ftell(fp) - 12) /
195                            (sizeof(int) +                /* framepos */
196                             sizeof(unsigned long long) + /* seek_pos */
197                             sizeof(unsigned long long) + /* seek_pos_dts */
198                             sizeof(unsigned long long)   /* pts */
199                            );
200
201         fseek(fp, 12, SEEK_SET);
202
203         idx->entries = MEM_callocN(sizeof(struct anim_index_entry) *
204                                    idx->num_entries, "anim_index_entries");
205
206         for (i = 0; i < idx->num_entries; i++) {
207                 fread(&idx->entries[i].frameno,
208                       sizeof(int), 1, fp);
209                 fread(&idx->entries[i].seek_pos,
210                       sizeof(unsigned long long), 1, fp);
211                 fread(&idx->entries[i].seek_pos_dts,
212                       sizeof(unsigned long long), 1, fp);
213                 fread(&idx->entries[i].pts,
214                       sizeof(unsigned long long), 1, fp);
215         }
216
217         if (((ENDIAN_ORDER == B_ENDIAN) != (header[8] == 'V'))) {
218                 for (i = 0; i < idx->num_entries; i++) {
219                         BLI_endian_switch_int32(&idx->entries[i].frameno);
220                         BLI_endian_switch_int64((int64_t *)&idx->entries[i].seek_pos);
221                         BLI_endian_switch_int64((int64_t *)&idx->entries[i].seek_pos_dts);
222                         BLI_endian_switch_int64((int64_t *)&idx->entries[i].pts);
223                 }
224         }
225
226         fclose(fp);
227
228         return idx;
229 }
230
231 unsigned long long IMB_indexer_get_seek_pos(
232         struct anim_index *idx, int frame_index)
233 {
234         if (frame_index < 0) {
235                 frame_index = 0;
236         }
237         if (frame_index >= idx->num_entries) {
238                 frame_index = idx->num_entries - 1;
239         }
240         return idx->entries[frame_index].seek_pos;
241 }
242
243 unsigned long long IMB_indexer_get_seek_pos_dts(
244         struct anim_index *idx, int frame_index)
245 {
246         if (frame_index < 0) {
247                 frame_index = 0;
248         }
249         if (frame_index >= idx->num_entries) {
250                 frame_index = idx->num_entries - 1;
251         }
252         return idx->entries[frame_index].seek_pos_dts;
253 }
254
255 int IMB_indexer_get_frame_index(struct anim_index *idx, int frameno)
256 {
257         int len = idx->num_entries;
258         int half;
259         int middle;
260         int first = 0;
261
262         /* bsearch (lower bound) the right index */
263         
264         while (len > 0) {
265                 half = len >> 1;
266                 middle = first;
267
268                 middle += half;
269
270                 if (idx->entries[middle].frameno < frameno) {
271                         first = middle;
272                         first++;
273                         len = len - half - 1;
274                 }
275                 else {
276                         len = half;
277                 }
278         }
279
280         if (first == idx->num_entries) {
281                 return idx->num_entries - 1;
282         }
283         else {
284                 return first;
285         }
286 }
287
288 unsigned long long IMB_indexer_get_pts(struct anim_index *idx,
289                                        int frame_index)
290 {
291         if (frame_index < 0) {
292                 frame_index = 0;
293         }
294         if (frame_index >= idx->num_entries) {
295                 frame_index = idx->num_entries - 1;
296         }
297         return idx->entries[frame_index].pts;
298 }
299
300 int IMB_indexer_get_duration(struct anim_index *idx)
301 {
302         if (idx->num_entries == 0) {
303                 return 0;
304         }
305         return idx->entries[idx->num_entries - 1].frameno + 1;
306 }
307
308 int IMB_indexer_can_scan(struct anim_index *idx,
309                          int old_frame_index, int new_frame_index)
310 {
311         /* makes only sense, if it is the same I-Frame and we are not
312          * trying to run backwards in time... */
313         return (IMB_indexer_get_seek_pos(idx, old_frame_index) == IMB_indexer_get_seek_pos(idx, new_frame_index) &&
314                 old_frame_index < new_frame_index);
315 }
316
317 void IMB_indexer_close(struct anim_index *idx)
318 {
319         MEM_freeN(idx->entries);
320         MEM_freeN(idx);
321 }
322
323 int IMB_proxy_size_to_array_index(IMB_Proxy_Size pr_size)
324 {
325         switch (pr_size) {
326                 case IMB_PROXY_NONE: /* if we got here, something is broken anyways,
327                                       * so sane defaults... */
328                         return 0;
329                 case IMB_PROXY_25:
330                         return 0;
331                 case IMB_PROXY_50:
332                         return 1;
333                 case IMB_PROXY_75:
334                         return 2;
335                 case IMB_PROXY_100:
336                         return 3;
337                 default:
338                         return 0;
339         }
340         return 0;
341 }
342
343 int IMB_timecode_to_array_index(IMB_Timecode_Type tc)
344 {
345         switch (tc) {
346                 case IMB_TC_NONE: /* if we got here, something is broken anyways,
347                                    * so sane defaults... */
348                         return 0;
349                 case IMB_TC_RECORD_RUN:
350                         return 0;
351                 case IMB_TC_FREE_RUN:
352                         return 1;
353                 case IMB_TC_INTERPOLATED_REC_DATE_FREE_RUN:
354                         return 2;
355                 case IMB_TC_RECORD_RUN_NO_GAPS:
356                         return 3;
357                 default:
358                         return 0;
359         }
360         return 0;
361 }
362
363
364 /* ----------------------------------------------------------------------
365  * - rebuild helper functions
366  * ---------------------------------------------------------------------- */
367
368 static void get_index_dir(struct anim *anim, char *index_dir, size_t index_dir_len)
369 {
370         if (!anim->index_dir[0]) {
371                 char fname[FILE_MAXFILE];
372                 BLI_split_dirfile(anim->name, index_dir, fname, index_dir_len, sizeof(fname));
373                 BLI_path_append(index_dir, index_dir_len, "BL_proxy");
374                 BLI_path_append(index_dir, index_dir_len, fname);
375         }
376         else {
377                 BLI_strncpy(index_dir, anim->index_dir, index_dir_len);
378         }
379 }
380
381 static void get_proxy_filename(struct anim *anim, IMB_Proxy_Size preview_size,
382                                char *fname, int temp)
383 {
384         char index_dir[FILE_MAXDIR];
385         int i = IMB_proxy_size_to_array_index(preview_size);
386
387         char proxy_name[256];
388         char proxy_temp_name[256];
389         char stream_suffix[20];
390         
391         stream_suffix[0] = 0;
392
393         if (anim->streamindex > 0) {
394                 BLI_snprintf(stream_suffix, sizeof(stream_suffix), "_st%d", anim->streamindex);
395         }
396
397         BLI_snprintf(proxy_name, sizeof(proxy_name), "proxy_%d%s.avi",
398                      (int) (proxy_fac[i] * 100), stream_suffix);
399         BLI_snprintf(proxy_temp_name, sizeof(proxy_temp_name), "proxy_%d%s_part.avi",
400                      (int) (proxy_fac[i] * 100), stream_suffix);
401
402         get_index_dir(anim, index_dir, sizeof(index_dir));
403
404         BLI_join_dirfile(fname, FILE_MAXFILE + FILE_MAXDIR, index_dir, 
405                          temp ? proxy_temp_name : proxy_name);
406 }
407
408 static void get_tc_filename(struct anim *anim, IMB_Timecode_Type tc,
409                             char *fname)
410 {
411         char index_dir[FILE_MAXDIR];
412         int i = IMB_timecode_to_array_index(tc);
413         const char *index_names[] = {
414                 "record_run%s.blen_tc",
415                 "free_run%s.blen_tc",
416                 "interp_free_run%s.blen_tc",
417                 "record_run_no_gaps%s.blen_tc"
418         };
419
420         char stream_suffix[20];
421         char index_name[256];
422         
423         stream_suffix[0] = 0;
424
425         if (anim->streamindex > 0) {
426                 BLI_snprintf(stream_suffix, 20, "_st%d", anim->streamindex);
427         }
428         
429         BLI_snprintf(index_name, 256, index_names[i], stream_suffix);
430
431         get_index_dir(anim, index_dir, sizeof(index_dir));
432         
433         BLI_join_dirfile(fname, FILE_MAXFILE + FILE_MAXDIR, 
434                          index_dir, index_name);
435 }
436
437 /* ----------------------------------------------------------------------
438  * - common rebuilder structures
439  * ---------------------------------------------------------------------- */
440
441 typedef struct IndexBuildContext {
442         int anim_type;
443 } IndexBuildContext;
444
445
446 /* ----------------------------------------------------------------------
447  * - ffmpeg rebuilder
448  * ---------------------------------------------------------------------- */
449
450 #ifdef WITH_FFMPEG
451
452 struct proxy_output_ctx {
453         AVFormatContext *of;
454         AVStream *st;
455         AVCodecContext *c;
456         AVCodec *codec;
457         struct SwsContext *sws_ctx;
458         AVFrame *frame;
459         uint8_t *video_buffer;
460         int video_buffersize;
461         int cfra;
462         int proxy_size;
463         int orig_height;
464         struct anim *anim;
465 };
466
467 // work around stupid swscaler 16 bytes alignment bug...
468
469 static int round_up(int x, int mod)
470 {
471         return x + ((mod - (x % mod)) % mod);
472 }
473
474 static struct proxy_output_ctx *alloc_proxy_output_ffmpeg(
475         struct anim *anim,
476         AVStream *st, int proxy_size, int width, int height,
477         int quality)
478 {
479         struct proxy_output_ctx *rv = MEM_callocN(
480                 sizeof(struct proxy_output_ctx), "alloc_proxy_output");
481         
482         char fname[FILE_MAX];
483         int ffmpeg_quality;
484
485         /* JPEG requires this */
486         width = round_up(width, 8);
487         height = round_up(height, 8);
488
489         rv->proxy_size = proxy_size;
490         rv->anim = anim;
491
492         get_proxy_filename(rv->anim, rv->proxy_size, fname, TRUE);
493         BLI_make_existing_file(fname);
494
495         rv->of = avformat_alloc_context();
496         rv->of->oformat = av_guess_format("avi", NULL, NULL);
497         
498         BLI_strncpy(rv->of->filename, fname, sizeof(rv->of->filename));
499
500         fprintf(stderr, "Starting work on proxy: %s\n", rv->of->filename);
501
502         rv->st = avformat_new_stream(rv->of, NULL);
503         rv->st->id = 0;
504
505         rv->c = rv->st->codec;
506         rv->c->codec_type = AVMEDIA_TYPE_VIDEO;
507         rv->c->codec_id = CODEC_ID_MJPEG;
508         rv->c->width = width;
509         rv->c->height = height;
510
511         rv->of->oformat->video_codec = rv->c->codec_id;
512         rv->codec = avcodec_find_encoder(rv->c->codec_id);
513
514         if (!rv->codec) {
515                 fprintf(stderr, "No ffmpeg MJPEG encoder available? "
516                         "Proxy not built!\n");
517                 av_free(rv->of);
518                 return NULL;
519         }
520
521         if (rv->codec->pix_fmts) {
522                 rv->c->pix_fmt = rv->codec->pix_fmts[0];
523         }
524         else {
525                 rv->c->pix_fmt = PIX_FMT_YUVJ420P;
526         }
527
528         rv->c->sample_aspect_ratio =
529             rv->st->sample_aspect_ratio =
530                 st->codec->sample_aspect_ratio;
531
532         rv->c->time_base.den = 25;
533         rv->c->time_base.num = 1;
534         rv->st->time_base = rv->c->time_base;
535
536         /* there's no  way to set JPEG quality in the same way as in AVI JPEG and image sequence,
537          * but this seems to be giving expected quality result */
538         ffmpeg_quality = (int)(1.0f + 30.0f * (1.0f - (float)quality / 100.0f) + 0.5f);
539         av_opt_set_int(rv->c, "qmin", ffmpeg_quality, 0);
540         av_opt_set_int(rv->c, "qmax", ffmpeg_quality, 0);
541
542         if (rv->of->flags & AVFMT_GLOBALHEADER) {
543                 rv->c->flags |= CODEC_FLAG_GLOBAL_HEADER;
544         }
545
546         if (avio_open(&rv->of->pb, fname, AVIO_FLAG_WRITE) < 0) {
547                 fprintf(stderr, "Couldn't open outputfile! "
548                         "Proxy not built!\n");
549                 av_free(rv->of);
550                 return 0;
551         }
552
553         avcodec_open2(rv->c, rv->codec, NULL);
554
555         rv->video_buffersize = 2000000;
556         rv->video_buffer = (uint8_t *)MEM_mallocN(
557                 rv->video_buffersize, "FFMPEG video buffer");
558
559         rv->orig_height = av_get_cropped_height_from_codec(st->codec);
560
561         if (st->codec->width != width || st->codec->height != height ||
562             st->codec->pix_fmt != rv->c->pix_fmt)
563         {
564                 rv->frame = avcodec_alloc_frame();
565                 avpicture_fill((AVPicture *) rv->frame,
566                                MEM_mallocN(avpicture_get_size(
567                                                rv->c->pix_fmt,
568                                                round_up(width, 16), height),
569                                            "alloc proxy output frame"),
570                                rv->c->pix_fmt, round_up(width, 16), height);
571
572                 rv->sws_ctx = sws_getContext(
573                         st->codec->width,
574                         rv->orig_height,
575                         st->codec->pix_fmt,
576                         width, height,
577                         rv->c->pix_fmt,
578                         SWS_FAST_BILINEAR | SWS_PRINT_INFO,
579                         NULL, NULL, NULL);
580         }
581
582         if (avformat_write_header(rv->of, NULL) < 0) {
583                 fprintf(stderr, "Couldn't set output parameters? "
584                         "Proxy not built!\n");
585                 av_free(rv->of);
586                 return 0;
587         }
588
589         return rv;
590 }
591
592 static int add_to_proxy_output_ffmpeg(
593         struct proxy_output_ctx *ctx, AVFrame *frame)
594 {
595         int outsize = 0;
596
597         if (!ctx) {
598                 return 0;
599         }
600
601         if (ctx->sws_ctx && frame &&
602             (frame->data[0] || frame->data[1] ||
603              frame->data[2] || frame->data[3]))
604         {
605                 sws_scale(ctx->sws_ctx, (const uint8_t *const *) frame->data,
606                           frame->linesize, 0, ctx->orig_height,
607                           ctx->frame->data, ctx->frame->linesize);
608         }
609
610         frame = ctx->sws_ctx ? (frame ? ctx->frame : 0) : frame;
611
612         if (frame) {
613                 frame->pts = ctx->cfra++;
614         }
615
616         outsize = avcodec_encode_video(
617                 ctx->c, ctx->video_buffer, ctx->video_buffersize,
618                 frame);
619
620         if (outsize < 0) {
621                 fprintf(stderr, "Error encoding proxy frame %d for '%s'\n", 
622                         ctx->cfra - 1, ctx->of->filename);
623                 return 0;
624         }
625
626         if (outsize != 0) {
627                 AVPacket packet;
628                 av_init_packet(&packet);
629
630                 if (ctx->c->coded_frame->pts != AV_NOPTS_VALUE) {
631                         packet.pts = av_rescale_q(ctx->c->coded_frame->pts,
632                                                   ctx->c->time_base,
633                                                   ctx->st->time_base);
634                 }
635                 if (ctx->c->coded_frame->key_frame)
636                         packet.flags |= AV_PKT_FLAG_KEY;
637
638                 packet.stream_index = ctx->st->index;
639                 packet.data = ctx->video_buffer;
640                 packet.size = outsize;
641
642                 if (av_interleaved_write_frame(ctx->of, &packet) != 0) {
643                         fprintf(stderr, "Error writing proxy frame %d "
644                                 "into '%s'\n", ctx->cfra - 1,
645                                 ctx->of->filename);
646                         return 0;
647                 }
648
649                 return 1;
650         }
651         else {
652                 return 0;
653         }
654 }
655
656 static void free_proxy_output_ffmpeg(struct proxy_output_ctx *ctx,
657                                      int rollback)
658 {
659         char fname[FILE_MAX];
660         char fname_tmp[FILE_MAX];
661
662         if (!ctx) {
663                 return;
664         }
665
666         if (!rollback) {
667                 while (add_to_proxy_output_ffmpeg(ctx, NULL)) {}
668         }
669
670         avcodec_flush_buffers(ctx->c);
671
672         av_write_trailer(ctx->of);
673         
674         avcodec_close(ctx->c);
675         
676         if (ctx->of->oformat) {
677                 if (!(ctx->of->oformat->flags & AVFMT_NOFILE)) {
678                         avio_close(ctx->of->pb);
679                 }
680         }
681         avformat_free_context(ctx->of);
682
683         MEM_freeN(ctx->video_buffer);
684
685         if (ctx->sws_ctx) {
686                 sws_freeContext(ctx->sws_ctx);
687
688                 MEM_freeN(ctx->frame->data[0]);
689                 av_free(ctx->frame);
690         }
691
692         get_proxy_filename(ctx->anim, ctx->proxy_size, 
693                            fname_tmp, TRUE);
694
695         if (rollback) {
696                 unlink(fname_tmp);
697         }
698         else {
699                 get_proxy_filename(ctx->anim, ctx->proxy_size, 
700                                    fname, FALSE);
701                 unlink(fname);
702                 BLI_rename(fname_tmp, fname);
703         }
704         
705         MEM_freeN(ctx);
706 }
707
708 typedef struct FFmpegIndexBuilderContext {
709         int anim_type;
710
711         AVFormatContext *iFormatCtx;
712         AVCodecContext *iCodecCtx;
713         AVCodec *iCodec;
714         AVStream *iStream;
715         int videoStream;
716
717         int num_proxy_sizes;
718         int num_indexers;
719
720         struct proxy_output_ctx *proxy_ctx[IMB_PROXY_MAX_SLOT];
721         anim_index_builder *indexer[IMB_TC_MAX_SLOT];
722
723         IMB_Timecode_Type tcs_in_use;
724         IMB_Proxy_Size proxy_sizes_in_use;
725
726         unsigned long long seek_pos;
727         unsigned long long last_seek_pos;
728         unsigned long long seek_pos_dts;
729         unsigned long long seek_pos_pts;
730         unsigned long long last_seek_pos_dts;
731         unsigned long long start_pts;
732         double frame_rate;
733         double pts_time_base;
734         int frameno, frameno_gapless;
735         int start_pts_set;
736 } FFmpegIndexBuilderContext;
737
738 static IndexBuildContext *index_ffmpeg_create_context(struct anim *anim, IMB_Timecode_Type tcs_in_use,
739                                                       IMB_Proxy_Size proxy_sizes_in_use, int quality)
740 {
741         FFmpegIndexBuilderContext *context = MEM_callocN(sizeof(FFmpegIndexBuilderContext), "FFmpeg index builder context");
742         int num_proxy_sizes = IMB_PROXY_MAX_SLOT;
743         int num_indexers = IMB_TC_MAX_SLOT;
744         int i, streamcount;
745
746         context->tcs_in_use = tcs_in_use;
747         context->proxy_sizes_in_use = proxy_sizes_in_use;
748         context->num_proxy_sizes = IMB_PROXY_MAX_SLOT;
749         context->num_indexers = IMB_TC_MAX_SLOT;
750
751         memset(context->proxy_ctx, 0, sizeof(context->proxy_ctx));
752         memset(context->indexer, 0, sizeof(context->indexer));
753
754         if (avformat_open_input(&context->iFormatCtx, anim->name, NULL, NULL) != 0) {
755                 MEM_freeN(context);
756                 return NULL;
757         }
758
759         if (avformat_find_stream_info(context->iFormatCtx, NULL) < 0) {
760                 av_close_input_file(context->iFormatCtx);
761                 MEM_freeN(context);
762                 return NULL;
763         }
764
765         streamcount = anim->streamindex;
766
767         /* Find the video stream */
768         context->videoStream = -1;
769         for (i = 0; i < context->iFormatCtx->nb_streams; i++)
770                 if (context->iFormatCtx->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO) {
771                         if (streamcount > 0) {
772                                 streamcount--;
773                                 continue;
774                         }
775                         context->videoStream = i;
776                         break;
777                 }
778
779         if (context->videoStream == -1) {
780                 av_close_input_file(context->iFormatCtx);
781                 MEM_freeN(context);
782                 return NULL;
783         }
784
785         context->iStream = context->iFormatCtx->streams[context->videoStream];
786         context->iCodecCtx = context->iStream->codec;
787
788         context->iCodec = avcodec_find_decoder(context->iCodecCtx->codec_id);
789
790         if (context->iCodec == NULL) {
791                 av_close_input_file(context->iFormatCtx);
792                 MEM_freeN(context);
793                 return NULL;
794         }
795
796         context->iCodecCtx->workaround_bugs = 1;
797
798         if (avcodec_open2(context->iCodecCtx, context->iCodec, NULL) < 0) {
799                 av_close_input_file(context->iFormatCtx);
800                 MEM_freeN(context);
801                 return NULL;
802         }
803
804         for (i = 0; i < num_proxy_sizes; i++) {
805                 if (proxy_sizes_in_use & proxy_sizes[i]) {
806                         context->proxy_ctx[i] = alloc_proxy_output_ffmpeg(
807                                 anim, context->iStream, proxy_sizes[i],
808                                 context->iCodecCtx->width * proxy_fac[i],
809                                 av_get_cropped_height_from_codec(
810                                 context->iCodecCtx) * proxy_fac[i],
811                                 quality);
812                         if (!context->proxy_ctx[i]) {
813                                 proxy_sizes_in_use &= ~proxy_sizes[i];
814                         }
815                 }
816         }
817
818         for (i = 0; i < num_indexers; i++) {
819                 if (tcs_in_use & tc_types[i]) {
820                         char fname[FILE_MAX];
821
822                         get_tc_filename(anim, tc_types[i], fname);
823
824                         context->indexer[i] = IMB_index_builder_create(fname);
825                         if (!context->indexer[i]) {
826                                 tcs_in_use &= ~tc_types[i];
827                         }
828                 }
829         }
830
831         return (IndexBuildContext *)context;
832 }
833
834 static void index_rebuild_ffmpeg_finish(FFmpegIndexBuilderContext *context, int stop)
835 {
836         int i;
837
838         for (i = 0; i < context->num_indexers; i++) {
839                 if (context->tcs_in_use & tc_types[i]) {
840                         IMB_index_builder_finish(context->indexer[i], stop);
841                 }
842         }
843
844         for (i = 0; i < context->num_proxy_sizes; i++) {
845                 if (context->proxy_sizes_in_use & proxy_sizes[i]) {
846                         free_proxy_output_ffmpeg(context->proxy_ctx[i], stop);
847                 }
848         }
849
850         avcodec_close(context->iCodecCtx);
851         avformat_close_input(&context->iFormatCtx);
852
853         MEM_freeN(context);
854 }
855
856 static void index_rebuild_ffmpeg_proc_decoded_frame(
857         FFmpegIndexBuilderContext *context, 
858         AVPacket * curr_packet,
859         AVFrame *in_frame)
860 {
861         int i;
862         unsigned long long s_pos = context->seek_pos;
863         unsigned long long s_dts = context->seek_pos_dts;
864         unsigned long long pts = av_get_pts_from_frame(context->iFormatCtx, in_frame);
865
866         for (i = 0; i < context->num_proxy_sizes; i++) {
867                 add_to_proxy_output_ffmpeg(context->proxy_ctx[i], in_frame);
868         }
869
870         if (!context->start_pts_set) {
871                 context->start_pts = pts;
872                 context->start_pts_set = TRUE;
873         }
874
875         context->frameno = floor((pts - context->start_pts) *
876                                  context->pts_time_base  *
877                                  context->frame_rate + 0.5);
878
879         /* decoding starts *always* on I-Frames,
880          * so: P-Frames won't work, even if all the
881          * information is in place, when we seek
882          * to the I-Frame presented *after* the P-Frame,
883          * but located before the P-Frame within
884          * the stream */
885
886         if (pts < context->seek_pos_pts) {
887                 s_pos = context->last_seek_pos;
888                 s_dts = context->last_seek_pos_dts;
889         }
890
891         for (i = 0; i < context->num_indexers; i++) {
892                 if (context->tcs_in_use & tc_types[i]) {
893                         int tc_frameno = context->frameno;
894
895                         if (tc_types[i] == IMB_TC_RECORD_RUN_NO_GAPS)
896                                 tc_frameno = context->frameno_gapless;
897                         
898                         IMB_index_builder_proc_frame(
899                                 context->indexer[i],
900                                 curr_packet->data,
901                                 curr_packet->size,
902                                 tc_frameno,
903                                 s_pos, s_dts, pts);
904                 }
905         }
906         
907         context->frameno_gapless++;
908 }
909
910 static int index_rebuild_ffmpeg(FFmpegIndexBuilderContext *context,
911                                 short *stop, short *do_update, float *progress)
912 {
913         AVFrame *in_frame = 0;
914         AVPacket next_packet;
915         uint64_t stream_size;
916
917         memset(&next_packet, 0, sizeof(AVPacket));
918
919         in_frame = avcodec_alloc_frame();
920
921         stream_size = avio_size(context->iFormatCtx->pb);
922
923         context->frame_rate = av_q2d(context->iStream->r_frame_rate);
924         context->pts_time_base = av_q2d(context->iStream->time_base);
925
926         while (av_read_frame(context->iFormatCtx, &next_packet) >= 0) {
927                 int frame_finished = 0;
928                 float next_progress =  (float)((int)floor(((double) next_packet.pos) * 100 /
929                                                           ((double) stream_size) + 0.5)) / 100;
930
931                 if (*progress != next_progress) {
932                         *progress = next_progress;
933                         *do_update = TRUE;
934                 }
935
936                 if (*stop) {
937                         av_free_packet(&next_packet);
938                         break;
939                 }
940
941                 if (next_packet.stream_index == context->videoStream) {
942                         if (next_packet.flags & AV_PKT_FLAG_KEY) {
943                                 context->last_seek_pos = context->seek_pos;
944                                 context->last_seek_pos_dts = context->seek_pos_dts;
945                                 context->seek_pos = next_packet.pos;
946                                 context->seek_pos_dts = next_packet.dts;
947                                 context->seek_pos_pts = next_packet.pts;
948                         }
949
950                         avcodec_decode_video2(
951                                 context->iCodecCtx, in_frame, &frame_finished,
952                                 &next_packet);
953                 }
954
955                 if (frame_finished) {
956                         index_rebuild_ffmpeg_proc_decoded_frame(
957                                 context, &next_packet, in_frame);
958                 }
959                 av_free_packet(&next_packet);
960         }
961
962         /* process pictures still stuck in decoder engine after EOF
963          * according to ffmpeg docs using 0-size packets.
964          *
965          * At least, if we haven't already stopped... */
966
967         /* this creates the 0-size packet and prevents a memory leak. */
968         av_free_packet(&next_packet);
969
970         if (!*stop) {
971                 int frame_finished;
972
973                 do {
974                         frame_finished = 0;
975
976                         avcodec_decode_video2(
977                                 context->iCodecCtx, in_frame, &frame_finished,
978                                 &next_packet);
979
980                         if (frame_finished) {
981                                 index_rebuild_ffmpeg_proc_decoded_frame(
982                                         context, &next_packet, in_frame);
983                         }
984                 } while (frame_finished);
985         }
986
987         av_free(in_frame);
988
989         return 1;
990 }
991
992 #endif
993
994 /* ----------------------------------------------------------------------
995  * - internal AVI (fallback) rebuilder
996  * ---------------------------------------------------------------------- */
997
998 #ifdef WITH_AVI
999 typedef struct FallbackIndexBuilderContext {
1000         int anim_type;
1001
1002         struct anim *anim;
1003         AviMovie *proxy_ctx[IMB_PROXY_MAX_SLOT];
1004         IMB_Proxy_Size proxy_sizes_in_use;
1005 } FallbackIndexBuilderContext;
1006
1007 static AviMovie *alloc_proxy_output_avi(
1008         struct anim *anim, char *filename, int width, int height,
1009         int quality)
1010 {
1011         int x, y;
1012         AviFormat format;
1013         double framerate;
1014         AviMovie *avi;
1015         short frs_sec = 25;      /* it doesn't really matter for proxies,
1016                                   * but sane defaults help anyways...*/
1017         float frs_sec_base = 1.0;
1018
1019         IMB_anim_get_fps(anim, &frs_sec, &frs_sec_base);
1020         
1021         x = width;
1022         y = height;
1023
1024         framerate = (double) frs_sec / (double) frs_sec_base;
1025         
1026         avi = MEM_mallocN(sizeof(AviMovie), "avimovie");
1027
1028         format = AVI_FORMAT_MJPEG;
1029
1030         if (AVI_open_compress(filename, avi, 1, format) != AVI_ERROR_NONE) {
1031                 MEM_freeN(avi);
1032                 return NULL;
1033         }
1034                         
1035         AVI_set_compress_option(avi, AVI_OPTION_TYPE_MAIN, 0, AVI_OPTION_WIDTH, &x);
1036         AVI_set_compress_option(avi, AVI_OPTION_TYPE_MAIN, 0, AVI_OPTION_HEIGHT, &y);
1037         AVI_set_compress_option(avi, AVI_OPTION_TYPE_MAIN, 0, AVI_OPTION_QUALITY, &quality);
1038         AVI_set_compress_option(avi, AVI_OPTION_TYPE_MAIN, 0, AVI_OPTION_FRAMERATE, &framerate);
1039
1040         avi->interlace = 0;
1041         avi->odd_fields = 0;
1042
1043         return avi;
1044 }
1045
1046 static IndexBuildContext *index_fallback_create_context(struct anim *anim, IMB_Timecode_Type UNUSED(tcs_in_use),
1047                                                         IMB_Proxy_Size proxy_sizes_in_use, int quality)
1048 {
1049         FallbackIndexBuilderContext *context;
1050         int i;
1051
1052         /* since timecode indices only work with ffmpeg right now,
1053          * don't know a sensible fallback here...
1054          *
1055          * so no proxies, no game to play...
1056          */
1057         if (proxy_sizes_in_use == IMB_PROXY_NONE) {
1058                 return NULL;
1059         }
1060
1061         context = MEM_callocN(sizeof(FallbackIndexBuilderContext), "fallback index builder context");
1062
1063         context->anim = anim;
1064         context->proxy_sizes_in_use = proxy_sizes_in_use;
1065
1066         memset(context->proxy_ctx, 0, sizeof(context->proxy_ctx));
1067
1068         for (i = 0; i < IMB_PROXY_MAX_SLOT; i++) {
1069                 if (context->proxy_sizes_in_use & proxy_sizes[i]) {
1070                         char fname[FILE_MAX];
1071
1072                         get_proxy_filename(anim, proxy_sizes[i], fname, TRUE);
1073                         BLI_make_existing_file(fname);
1074
1075                         context->proxy_ctx[i] = alloc_proxy_output_avi(anim, fname,
1076                                                                        anim->x * proxy_fac[i], anim->y * proxy_fac[i], quality);
1077                 }
1078         }
1079
1080         return (IndexBuildContext *)context;
1081 }
1082
1083 static void index_rebuild_fallback_finish(FallbackIndexBuilderContext *context, int stop)
1084 {
1085         struct anim *anim = context->anim;
1086         char fname[FILE_MAX];
1087         char fname_tmp[FILE_MAX];
1088         int i;
1089
1090         for (i = 0; i < IMB_PROXY_MAX_SLOT; i++) {
1091                 if (context->proxy_sizes_in_use & proxy_sizes[i]) {
1092                         AVI_close_compress(context->proxy_ctx[i]);
1093                         MEM_freeN(context->proxy_ctx[i]);
1094
1095                         get_proxy_filename(anim, proxy_sizes[i], fname_tmp, TRUE);
1096                         get_proxy_filename(anim, proxy_sizes[i], fname, FALSE);
1097
1098                         if (stop) {
1099                                 unlink(fname_tmp);
1100                         }
1101                         else {
1102                                 unlink(fname);
1103                                 rename(fname_tmp, fname);
1104                         }
1105                 }
1106         }
1107 }
1108
1109 static void index_rebuild_fallback(FallbackIndexBuilderContext *context,
1110                                    short *stop, short *do_update, float *progress)
1111 {
1112         int cnt = IMB_anim_get_duration(context->anim, IMB_TC_NONE);
1113         int i, pos;
1114         struct anim *anim = context->anim;
1115
1116         for (pos = 0; pos < cnt; pos++) {
1117                 struct ImBuf *ibuf = IMB_anim_absolute(anim, pos, IMB_TC_NONE, IMB_PROXY_NONE);
1118                 struct ImBuf *tmp_ibuf = IMB_dupImBuf(ibuf);
1119                 float next_progress = (float) pos / (float) cnt;
1120
1121                 if (*progress != next_progress) {
1122                         *progress = next_progress;
1123                         *do_update = TRUE;
1124                 }
1125                 
1126                 if (*stop) {
1127                         break;
1128                 }
1129
1130                 IMB_flipy(tmp_ibuf);
1131
1132                 for (i = 0; i < IMB_PROXY_MAX_SLOT; i++) {
1133                         if (context->proxy_sizes_in_use & proxy_sizes[i]) {
1134                                 int x = anim->x * proxy_fac[i];
1135                                 int y = anim->y * proxy_fac[i];
1136
1137                                 struct ImBuf *s_ibuf = IMB_dupImBuf(tmp_ibuf);
1138
1139                                 IMB_scalefastImBuf(s_ibuf, x, y);
1140
1141                                 IMB_convert_rgba_to_abgr(s_ibuf);
1142         
1143                                 AVI_write_frame(context->proxy_ctx[i], pos,
1144                                                 AVI_FORMAT_RGB32,
1145                                                 s_ibuf->rect, x * y * 4);
1146
1147                                 /* note that libavi free's the buffer... */
1148                                 s_ibuf->rect = NULL;
1149
1150                                 IMB_freeImBuf(s_ibuf);
1151                         }
1152                 }
1153
1154                 IMB_freeImBuf(tmp_ibuf);
1155                 IMB_freeImBuf(ibuf);
1156         }
1157 }
1158
1159 #endif  /* WITH_AVI */
1160
1161 /* ----------------------------------------------------------------------
1162  * - public API
1163  * ---------------------------------------------------------------------- */
1164
1165 IndexBuildContext *IMB_anim_index_rebuild_context(struct anim *anim, IMB_Timecode_Type tcs_in_use,
1166                                                   IMB_Proxy_Size proxy_sizes_in_use, int quality)
1167 {
1168         IndexBuildContext *context = NULL;
1169
1170         switch (anim->curtype) {
1171 #ifdef WITH_FFMPEG
1172                 case ANIM_FFMPEG:
1173                         context = index_ffmpeg_create_context(anim, tcs_in_use, proxy_sizes_in_use, quality);
1174                         break;
1175 #endif
1176 #ifdef WITH_AVI
1177                 default:
1178                         context = index_fallback_create_context(anim, tcs_in_use, proxy_sizes_in_use, quality);
1179                         break;
1180 #endif
1181         }
1182
1183         if (context)
1184                 context->anim_type = anim->curtype;
1185
1186         return context;
1187
1188         (void)tcs_in_use, (void)proxy_sizes_in_use, (void)quality;
1189 }
1190
1191 void IMB_anim_index_rebuild(struct IndexBuildContext *context,
1192                             short *stop, short *do_update, float *progress)
1193 {
1194         switch (context->anim_type) {
1195 #ifdef WITH_FFMPEG
1196                 case ANIM_FFMPEG:
1197                         index_rebuild_ffmpeg((FFmpegIndexBuilderContext *)context, stop, do_update, progress);
1198                         break;
1199 #endif
1200 #ifdef WITH_AVI
1201                 default:
1202                         index_rebuild_fallback((FallbackIndexBuilderContext *)context, stop, do_update, progress);
1203                         break;
1204 #endif
1205         }
1206
1207         (void)stop, (void)do_update, (void)progress;
1208 }
1209
1210 void IMB_anim_index_rebuild_finish(IndexBuildContext *context, short stop)
1211 {
1212         switch (context->anim_type) {
1213 #ifdef WITH_FFMPEG
1214                 case ANIM_FFMPEG:
1215                         index_rebuild_ffmpeg_finish((FFmpegIndexBuilderContext *)context, stop);
1216                         break;
1217 #endif
1218 #ifdef WITH_AVI
1219                 default:
1220                         index_rebuild_fallback_finish((FallbackIndexBuilderContext *)context, stop);
1221                         break;
1222 #endif
1223         }
1224
1225         (void)stop;
1226         (void)proxy_sizes;  /* static defined at top of the file */
1227 }
1228
1229
1230 void IMB_free_indices(struct anim *anim)
1231 {
1232         int i;
1233
1234         for (i = 0; i < IMB_PROXY_MAX_SLOT; i++) {
1235                 if (anim->proxy_anim[i]) {
1236                         IMB_close_anim(anim->proxy_anim[i]);
1237                         anim->proxy_anim[i] = NULL;
1238                 }
1239         }
1240
1241         for (i = 0; i < IMB_TC_MAX_SLOT; i++) {
1242                 if (anim->curr_idx[i]) {
1243                         IMB_indexer_close(anim->curr_idx[i]);
1244                         anim->curr_idx[i] = NULL;
1245                 }
1246         }
1247
1248
1249         anim->proxies_tried = 0;
1250         anim->indices_tried = 0;
1251 }
1252
1253 void IMB_anim_set_index_dir(struct anim *anim, const char *dir)
1254 {
1255         if (strcmp(anim->index_dir, dir) == 0) {
1256                 return;
1257         }
1258         BLI_strncpy(anim->index_dir, dir, sizeof(anim->index_dir));
1259
1260         IMB_free_indices(anim);
1261 }
1262
1263 struct anim *IMB_anim_open_proxy(
1264         struct anim *anim, IMB_Proxy_Size preview_size)
1265 {
1266         char fname[FILE_MAX];
1267         int i = IMB_proxy_size_to_array_index(preview_size);
1268
1269         if (anim->proxy_anim[i]) {
1270                 return anim->proxy_anim[i];
1271         }
1272
1273         if (anim->proxies_tried & preview_size) {
1274                 return NULL;
1275         }
1276
1277         get_proxy_filename(anim, preview_size, fname, FALSE);
1278
1279         /* proxies are generated in default color space */
1280         anim->proxy_anim[i] = IMB_open_anim(fname, 0, 0, NULL);
1281         
1282         anim->proxies_tried |= preview_size;
1283
1284         return anim->proxy_anim[i];
1285 }
1286
1287 struct anim_index *IMB_anim_open_index(
1288         struct anim *anim, IMB_Timecode_Type tc)
1289 {
1290         char fname[FILE_MAX];
1291         int i = IMB_timecode_to_array_index(tc);
1292
1293         if (anim->curr_idx[i]) {
1294                 return anim->curr_idx[i];
1295         }
1296
1297         if (anim->indices_tried & tc) {
1298                 return NULL;
1299         }
1300
1301         get_tc_filename(anim, tc, fname);
1302
1303         anim->curr_idx[i] = IMB_indexer_open(fname);
1304         
1305         anim->indices_tried |= tc;
1306
1307         return anim->curr_idx[i];
1308 }
1309
1310 int IMB_anim_index_get_frame_index(struct anim *anim, IMB_Timecode_Type tc,
1311                                    int position)
1312 {
1313         struct anim_index *idx = IMB_anim_open_index(anim, tc);
1314
1315         if (!idx) {
1316                 return position;
1317         }
1318
1319         return IMB_indexer_get_frame_index(idx, position);
1320 }
1321