<commit_before>import numpy as np

import click
import json
import os
import skvideo.io
import tensorflow as tf

from PIL import Image, ImageDraw
from luminoth.utils.predicting import PredictorNetwork

IMAGE_FORMATS = ['jpg', 'jpeg', 'png']
VIDEO_FORMATS = ['mov', 'mp4']  # TODO: check if more formats work


def get_filetype(filename):
    extension = filename.split('.')[-1].lower()
    if extension in IMAGE_FORMATS:
        return 'image'
    elif extension in VIDEO_FORMATS:
        return 'video'


@click.command(help='Obtain a model\'s predictions on an image or directory of images.')  # noqa
@click.argument('path-or-dir')
@click.option('config_files', '--config', '-c', required=True, multiple=True, help='Config to use.')  # noqa
@click.option('--output-dir', help='Where to write output')
@click.option('--save/--no-save', default=False, help='Save the image with the prediction of the model')  # noqa
@click.option('--min-prob', default=0.5, type=float, help='When drawing, only draw bounding boxes with probability larger than.')  # noqa
@click.option('--ignore-classes', default=None, multiple=True, help='Classes to ignore when predicting')  # noqa
@click.option('--debug', is_flag=True, help='Set debug level logging.')
def predict(path_or_dir, config_files, output_dir, save, min_prob,
            ignore_classes, debug):
    if debug:
        tf.logging.set_verbosity(tf.logging.DEBUG)
    else:
        tf.logging.set_verbosity(tf.logging.INFO)

    # -- Get file paths --
    if tf.gfile.IsDirectory(path_or_dir):
        file_paths = [
            os.path.join(path_or_dir, f)
            for f in tf.gfile.ListDirectory(path_or_dir)
            if get_filetype(f) in ('image', 'video')
        ]
    else:
        if get_filetype(path_or_dir) in ('image', 'video'):
            file_paths = [path_or_dir]
        else:
            file_paths = []

    errors = 0
    successes = 0
    created_files_paths = []
    total_files = len(file_paths)
    if total_files == 0:
        no_files_message = ("No images or videos found. "
                            "Accepted formats -> Image: {} - Video: {}")
        tf.logging.error(no_files_message.format(IMAGE_FORMATS, VIDEO_FORMATS))
        exit()

    # -- Initialize model --
    network = PredictorNetwork(config_files)

    #  -- Create output_dir if it doesn't exist --
    if output_dir:
        tf.gfile.MakeDirs(output_dir)

    tf.logging.info('Getting predictions for {} files'.format(total_files))

    # -- Iterate over file paths --
    for file_path in file_paths:

        save_path = 'pred_' + os.path.basename(file_path)
        if output_dir:
            save_path = os.path.join(output_dir, save_path)

        if get_filetype(file_path) == 'image':
            print('Predicting {}...'.format(file_path))
            with tf.gfile.Open(file_path, 'rb') as f:
                try:
                    image = Image.open(f).convert('RGB')
                except tf.errors.OutOfRangeError as e:
                    tf.logging.warning('Error: {}'.format(e))
                    tf.logging.warning('{} failed.'.format(file_path))
                    errors += 1
                    continue

            # Run image through network
            prediction = network.predict_image(image)
            successes += 1

            # Filter results if required by user
            if ignore_classes:
                prediction = filter_classes(prediction, ignore_classes)

            # Save prediction json file
            with open(save_path + '.json', 'w') as outfile:
                json.dump(prediction, outfile)
            created_files_paths.append(save_path + '.json')

            # Save predicted image
            if save:
                with tf.gfile.Open(file_path, 'rb') as im_file:
                    image = Image.open(im_file)
                    draw_bboxes_on_image(image, prediction, min_prob)
                    image.save(save_path)
                created_files_paths.append(save_path)

        elif get_filetype(file_path) == 'video':
            # NOTE: We'll hardcode the video ouput to mp4 for the time being
            save_path = os.path.splitext(save_path)[0] + '.mp4'
            try:
                writer = skvideo.io.FFmpegWriter(save_path)
            except AssertionError as e:
                tf.logging.error(e)
                tf.logging.error(
                    "Please install ffmpeg before making video predictions."
                )
                exit()
            num_of_frames = int(
                skvideo.io.ffprobe(file_path)['video']['@nb_frames'])
            video_progress_bar = click.progressbar(
                skvideo.io.vreader(file_path),
                length=num_of_frames,
                label='Predicting {}'.format(file_path))
            with video_progress_bar as bar:
                for frame in bar:
                    # Run image through network
                    prediction = network.predict_image(frame)

                    # Filter results if required by user
                    if ignore_classes:
                        prediction = filter_classes(prediction, ignore_classes)

                    image = Image.fromarray(frame)
                    draw_bboxes_on_image(image, prediction, min_prob)
                    writer.writeFrame(np.array(image))
                writer.close()
                created_files_paths.append(save_path)

        else:
            tf.logging.warning("{} isn't an image/video".format(file_path))

    # -- Generate logs --
    tf.logging.info(
        "Created the following files: {}".format(
            ', '.join(created_files_paths)
        )
    )<commit_msg>Handle corrupted videos in predict.py<commit_after>import numpy as np

import click
import json
import os
import skvideo.io
import tensorflow as tf

from PIL import Image, ImageDraw
from luminoth.utils.predicting import PredictorNetwork

IMAGE_FORMATS = ['jpg', 'jpeg', 'png']
VIDEO_FORMATS = ['mov', 'mp4', 'avi']  # TODO: check if more formats work


def get_filetype(filename):
    extension = filename.split('.')[-1].lower()
    if extension in IMAGE_FORMATS:
        return 'image'
    elif extension in VIDEO_FORMATS:
        return 'video'


@click.command(help='Obtain a model\'s predictions on an image or directory of images.')  # noqa
@click.argument('path-or-dir')
@click.option('config_files', '--config', '-c', required=True, multiple=True, help='Config to use.')  # noqa
@click.option('--output-dir', help='Where to write output')
@click.option('--save/--no-save', default=False, help='Save the image with the prediction of the model')  # noqa
@click.option('--min-prob', default=0.5, type=float, help='When drawing, only draw bounding boxes with probability larger than.')  # noqa
@click.option('--ignore-classes', default=None, multiple=True, help='Classes to ignore when predicting')  # noqa
@click.option('--debug', is_flag=True, help='Set debug level logging.')
def predict(path_or_dir, config_files, output_dir, save, min_prob,
            ignore_classes, debug):
    if debug:
        tf.logging.set_verbosity(tf.logging.DEBUG)
    else:
        tf.logging.set_verbosity(tf.logging.INFO)

    # -- Get file paths --
    if tf.gfile.IsDirectory(path_or_dir):
        file_paths = [
            os.path.join(path_or_dir, f)
            for f in tf.gfile.ListDirectory(path_or_dir)
            if get_filetype(f) in ('image', 'video')
        ]
    else:
        if get_filetype(path_or_dir) in ('image', 'video'):
            file_paths = [path_or_dir]
        else:
            file_paths = []

    errors = 0
    successes = 0
    created_files_paths = []
    total_files = len(file_paths)
    if total_files == 0:
        no_files_message = ("No images or videos found. "
                            "Accepted formats -> Image: {} - Video: {}")
        tf.logging.error(no_files_message.format(IMAGE_FORMATS, VIDEO_FORMATS))
        exit()

    # -- Initialize model --
    network = PredictorNetwork(config_files)

    #  -- Create output_dir if it doesn't exist --
    if output_dir:
        tf.gfile.MakeDirs(output_dir)

    tf.logging.info('Getting predictions for {} files'.format(total_files))

    # -- Iterate over file paths --
    for file_path in file_paths:

        save_path = 'pred_' + os.path.basename(file_path)
        if output_dir:
            save_path = os.path.join(output_dir, save_path)

        if get_filetype(file_path) == 'image':
            print('Predicting {}...'.format(file_path))
            with tf.gfile.Open(file_path, 'rb') as f:
                try:
                    image = Image.open(f).convert('RGB')
                except tf.errors.OutOfRangeError as e:
                    tf.logging.warning('Error: {}'.format(e))
                    tf.logging.warning('{} failed.'.format(file_path))
                    errors += 1
                    continue

            # Run image through network
            prediction = network.predict_image(image)
            successes += 1

            # Filter results if required by user
            if ignore_classes:
                prediction = filter_classes(prediction, ignore_classes)

            # Save prediction json file
            with open(save_path + '.json', 'w') as outfile:
                json.dump(prediction, outfile)
            created_files_paths.append(save_path + '.json')

            # Save predicted image
            if save:
                with tf.gfile.Open(file_path, 'rb') as im_file:
                    image = Image.open(im_file)
                    draw_bboxes_on_image(image, prediction, min_prob)
                    image.save(save_path)
                created_files_paths.append(save_path)

        elif get_filetype(file_path) == 'video':
            # NOTE: We'll hardcode the video ouput to mp4 for the time being
            save_path = os.path.splitext(save_path)[0] + '.mp4'
            try:
                writer = skvideo.io.FFmpegWriter(save_path)
            except AssertionError as e:
                tf.logging.error(e)
                tf.logging.error(
                    "Please install ffmpeg before making video predictions."
                )
                exit()
            num_of_frames = int(
                skvideo.io.ffprobe(file_path)['video']['@nb_frames'])
            video_progress_bar = click.progressbar(
                skvideo.io.vreader(file_path),
                length=num_of_frames,
                label='Predicting {}'.format(file_path))
            with video_progress_bar as bar:
                try:
                    for frame in bar:
                        # Run image through network
                        prediction = network.predict_image(frame)

                        # Filter results if required by user
                        if ignore_classes:
                            prediction = filter_classes(prediction, ignore_classes)

                        image = Image.fromarray(frame)
                        draw_bboxes_on_image(image, prediction, min_prob)
                        writer.writeFrame(np.array(image))
                except RuntimeError as e:
                    print()  # So error won't get printed next to progressbar
                    tf.logging.error('Error: {}'.format(e))
                    tf.logging.error('Corrupt videofile: {}'.format(file_path))
                    tf.logging.error(
                        'Partially processed video file saved in {}'.format(
                            save_path))
                    errors += 1

            writer.close()
            created_files_paths.append(save_path)

        else:
            tf.logging.warning("{} isn't an image/video".format(file_path))

    # -- Generate logs --
    tf.logging.info(
        "Created the following files: {}".format(
            ', '.join(created_files_paths)
        )
    )
{'fraction_non_alphanumeric': 0.10700454621990234, 'fraction_numerical': 0.002104731436268732, 'mean_word_length': 1.956445993031359, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1261643', 'n_tokens_mistral': 3422, 'n_tokens_neox': 3281, 'n_words': 801}