diff --git a/dockerfile b/dockerfile index c48b236..718986e 100644 --- a/dockerfile +++ b/dockerfile @@ -12,7 +12,7 @@ RUN apt-get update && apt-get install -y apt-utils software-properties-common && libavformat-dev libswscale-dev unzip && \ python3.6 -m pip install --upgrade pip && \ python3.6 -m pip install jupyter ipywidgets jupyterlab && \ - python3.6 -m pip install tensorflow-gpu h5py keras && \ + python3.6 -m pip install h5py keras && \ python3.6 -m pip install scikit-image opencv-contrib-python pyyaml RUN mkdir /comixify diff --git a/keyframes/keyframes.py b/keyframes/keyframes.py index 544a826..34c9983 100644 --- a/keyframes/keyframes.py +++ b/keyframes/keyframes.py @@ -61,7 +61,7 @@ class KeyFramesExtractor: jj(f"{settings.TMP_DIR}", f"{all_frames_tmp_dir}", "%06d.jpeg")]) else: call(["ffmpeg", "-i", video.file.path, "-vf", "select=not(mod(n\\,15))", "-vsync", "vfr", "-q:v", "2", - jj(settings.TMP_DIR, all_frames_tmp_dir, "%06d.jpeg")]) + jj(settings.TMP_DIR, all_frames_tmp_dir, "%06d.bmp")]) frames_paths = [] for dirname, dirnames, filenames in os.walk(jj(settings.TMP_DIR, all_frames_tmp_dir)): for filename in filenames: diff --git a/neural_image_assessment/models.py b/neural_image_assessment/models.py index fda9d51..ff3a81c 100644 --- a/neural_image_assessment/models.py +++ b/neural_image_assessment/models.py @@ -2,11 +2,12 @@ import os import errno import numpy as np from keras.models import load_model -from keras.preprocessing.image import load_img, img_to_array +from keras.preprocessing.image import img_to_array from keras.applications.nasnet import preprocess_input import tensorflow as tf from PIL import Image + MODEL_PATH = 'neural_image_assessment/pretrained_model/nima_model.h5' @@ -16,12 +17,11 @@ class NeuralImageAssessment: raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), MODEL_PATH) self.graph = tf.Graph() config = tf.ConfigProto() - config.gpu_options.per_process_gpu_memory_fraction = 0.001 + config.gpu_options.per_process_gpu_memory_fraction = 0.1 config.gpu_options.allow_growth = True self.session = tf.Session(graph=self.graph, config=config) with self.graph.as_default(): - with tf.device('/CPU:0'): - self.model = load_model(MODEL_PATH) + self.model = load_model(MODEL_PATH) @staticmethod def resize_image(bgr_img_array, target_size=(224, 224), interpolation='nearest'): diff --git a/requirements.txt b/requirements.txt index 2eeef33..7d4a461 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,3 +13,4 @@ torch==0.4.1 torchvision==0.2.1 scikit-learn==0.19.2 youtube-dl==2018.9.18 +tensorflow-gpu==1.10.1 \ No newline at end of file