diff --git a/sloth/annotations/container.py b/sloth/annotations/container.py index 62ccde8..85c93d3 100644 --- a/sloth/annotations/container.py +++ b/sloth/annotations/container.py @@ -1,9 +1,14 @@ import os +import sys import fnmatch import time +import numpy as np from sloth.core.exceptions import \ ImproperlyConfigured, NotImplementedException, InvalidArgumentException from sloth.core.utils import import_callable +import logging +LOG = logging.getLogger(__name__) + try: import cPickle as pickle except: @@ -16,9 +21,15 @@ try: import yaml except: pass -import okapy -import logging -LOG = logging.getLogger(__name__) +try: + import okapy + _use_pil = False +except: + try: + from PIL import Image + _use_pil = True + except: + LOG.warn("Could neither find PIL nor okapy. Sloth needs one of them for loading images.") class AnnotationContainerFactory: @@ -138,7 +149,15 @@ class AnnotationContainer: relative to the label file's directory. """ fullpath = self._fullpath(filename) - return okapy.loadImage(fullpath) + if not os.path.exists(fullpath): + LOG.warn("Image file %s does not exists." % fullpath) + return None + + if _use_pil: + im = Image.open(fullpath) + return np.asarray(im) + else: + return okapy.loadImage(fullpath) def loadFrame(self, filename, frame_number): """ @@ -147,6 +166,10 @@ class AnnotationContainer: the video from a path relative to the label files directory. """ fullpath = self._fullpath(filename) + if not os.path.exists(fullpath): + LOG.warn("Video file %s does not exists." % fullpath) + return None + #TODO load video diff --git a/sloth/core/labeltool.py b/sloth/core/labeltool.py index 3809da6..114ee6b 100755 --- a/sloth/core/labeltool.py +++ b/sloth/core/labeltool.py @@ -19,7 +19,10 @@ from sloth.gui import MainWindow import logging LOG = logging.getLogger(__name__) -import okapy.videoio as okv +try: + import okapy.videoio as okv +except: + pass class LabelTool(QObject): """ diff --git a/sloth/gui/annotationscene.py b/sloth/gui/annotationscene.py index b1faf72..c7ae31b 100644 --- a/sloth/gui/annotationscene.py +++ b/sloth/gui/annotationscene.py @@ -4,7 +4,7 @@ from PyQt4.QtCore import * from sloth.items import * from sloth.core.exceptions import InvalidArgumentException from sloth.annotations.model import AnnotationModelItem -import okapy +from sloth.utils import toQImage import logging LOG = logging.getLogger(__name__) @@ -74,7 +74,7 @@ class AnnotationScene(QGraphicsScene): self.image_item_ = current_image assert self.image_item_.model() == self.model_ self.image_ = self.labeltool_.getImage(self.image_item_) - self.pixmap_ = QPixmap(okapy.guiqt.toQImage(self.image_)) + self.pixmap_ = QPixmap(toQImage(self.image_)) item = QGraphicsPixmapItem(self.pixmap_) item.setZValue(-1) self.setSceneRect(0, 0, self.pixmap_.width(), self.pixmap_.height()) diff --git a/sloth/gui/frameviewer.py b/sloth/gui/frameviewer.py index 14d5097..c8676f5 100644 --- a/sloth/gui/frameviewer.py +++ b/sloth/gui/frameviewer.py @@ -1,10 +1,12 @@ #!/usr/bin/python import sys, os, math -import okapy -import okapy.videoio -import okapy.guiqt.utilities as ogu +from sloth.utils import toQImage from PyQt4.QtCore import * from PyQt4.QtGui import * +try: + import okapy.videoio +except: + pass videos = [] scenes = [] @@ -197,7 +199,7 @@ def get_dummy_scene(): video.getFrame(1000) video.getNextFrame() img = video.getImage() - qimg = ogu.toQImage(img, True) + qimg = toQImage(img, True) scene.addPixmap(QPixmap(qimg)) return (scene, video) @@ -206,7 +208,7 @@ def next_frame(): v.getNextFrame() for v, s in zip(videos, scenes): img = v.getImage() - qimg = ogu.toQImage(img, True) + qimg = toQImage(img, True) s.clear() s.addPixmap(QPixmap(qimg)) diff --git a/sloth/utils/__init__.py b/sloth/utils/__init__.py index e69de29..69eb24c 100644 --- a/sloth/utils/__init__.py +++ b/sloth/utils/__init__.py @@ -0,0 +1,25 @@ +from sloth.core.exceptions import NotImplementedException +from PyQt4.QtGui import QImage, qRgb +import numpy as np + +gray_color_table = [qRgb(i, i, i) for i in range(256)] +def toQImage(im, copy=False): + if im is None: + return QImage() + + if im.dtype == np.uint8: + if len(im.shape) == 2: + qim = QImage(im.data, im.shape[1], im.shape[0], im.strides[0], QImage.Format_Indexed8) + qim.setColorTable(gray_color_table) + return qim.copy() if copy else qim + + elif len(im.shape) == 3: + if im.shape[2] == 3: + qim = QImage(im.data, im.shape[1], im.shape[0], im.strides[0], QImage.Format_RGB888); + return qim.copy() if copy else qim + elif im.shape[2] == 4: + qim = QImage(im.data, im.shape[1], im.shape[0], im.strides[0], QImage.Format_ARGB32); + return qim.copy() if copy else qim + + raise NotImplementedException +