<commit_before># Copyright (C) 2010-2012 Cuckoo Sandbox Developers.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import sys
import json
from datetime import datetime

from lib.cuckoo.common.constants import CUCKOO_ROOT
from lib.cuckoo.common.exceptions import CuckooDatabaseError
from lib.cuckoo.common.exceptions import CuckooOperationalError
from lib.cuckoo.common.exceptions import CuckooDependencyError
from lib.cuckoo.common.config import Config
from lib.cuckoo.common.objects import File, URL
from lib.cuckoo.common.utils import create_folder

try:
    from sqlalchemy import create_engine, Column
    from sqlalchemy import Integer, String, Boolean, DateTime, Enum
    from sqlalchemy import ForeignKey, Text, Index
    from sqlalchemy.orm import sessionmaker, relationship
    from sqlalchemy.sql import func
    from sqlalchemy.ext.declarative import declarative_base
    from sqlalchemy.exc import SQLAlchemyError, IntegrityError
    from sqlalchemy.pool import NullPool
    Base = declarative_base()
except ImportError:
    raise CuckooDependencyError("SQLAlchemy library not found, "
                                "verify your setup")

class Task(Base):
    """Analysis task queue."""
    __tablename__ = "tasks"

    id = Column(Integer(), primary_key=True)
    target = Column(Text(), nullable=False)
    category = Column(String(255), nullable=False)
    timeout = Column(Integer(), server_default="0", nullable=False)
    priority = Column(Integer(), server_default="1", nullable=False)
    custom = Column(String(255), nullable=True)
    machine = Column(String(255), nullable=True)
    package = Column(String(255), nullable=True)
    options = Column(String(255), nullable=True)
    platform = Column(String(255), nullable=True)
    added_on = Column(DateTime(timezone=False),
                      default=datetime.now(),
                      nullable=False)
    completed_on = Column(DateTime(timezone=False), nullable=True)
    status = Column(Enum("pending",
                         "processing",
                         "failure",
                         "success",
                         name="status_type"),
                         server_default="pending",
                         nullable=False)
    sample_id = Column(Integer, ForeignKey("samples.id"), nullable=True)
    sample = relationship("Sample", backref="tasks")
    guest = relationship("Guest", uselist=False, backref="tasks")

    def to_dict(self):
        """Converts object to dict.
        @return: dict
        """
        d = {}
        for column in self.__table__.columns:
            value = getattr(self, column.name)
            if isinstance(value, datetime):
                d[column.name] = value.strftime("%Y-%m-%d %H:%M:%S")
            else:
                d[column.name] = value
        return d

    def to_json(self):
        """Converts object to JSON.
        @return: JSON data
        """
        return json.dumps(self.to_dict())

    def __init__(self, target=None):
        self.target = target

    def __repr__(self):
        return "<Task('%s','%s')>" % (self.id, self.target)

class Guest(Base):
    """Tracks guest run."""
    __tablename__ = "guests"

    id = Column(Integer(), primary_key=True)
    name = Column(String(255), nullable=False)
    label = Column(String(255), nullable=False)
    manager = Column(String(255), nullable=False)
    started_on = Column(DateTime(timezone=False),
                        default=datetime.now(),
                        nullable=False)
    shutdown_on = Column(DateTime(timezone=False), nullable=True)
    task_id = Column(Integer,
                     ForeignKey('tasks.id'),
                     nullable=False,
                     unique=True)

    def __repr__(self):
        return "<Guest('%s','%s')>" % (self.id, self.name)

    def to_dict(self):
        """Converts object to dict.
        @return: dict
        """
        d = {}
        for column in self.__table__.columns:
            value = getattr(self, column.name)
            if isinstance(value, datetime):
                d[column.name] = value.strftime("%Y-%m-%d %H:%M:%S")
            else:
                d[column.name] = value

    def to_json(self):
        """Converts object to JSON.
        @return: JSON data
        """
        return json.dumps(self.to_dict())

    def __init__(self, name, label, manager):
        self.name = name
        self.label = label
        self.manager = manager

class Sample(Base):
    """Submitted files details."""
    __tablename__ = "samples"

    id = Column(Integer(), primary_key=True)
    file_size = Column(Integer(), nullable=False)
    file_type = Column(String(255), nullable=False)
    md5 = Column(String(32), nullable=False)
    crc32 = Column(String(8), nullable=False)
    sha1 = Column(String(40), nullable=False)
    sha256 = Column(String(64), nullable=False)
    sha512 = Column(String(128), nullable=False)
    ssdeep = Column(String(255), nullable=True)
    __table_args__ = (Index("hash_index",
                            "md5",
                            "crc32",
                            "sha1",
                            "sha256",
                            "sha512",
                            unique=True), )

    def __repr__(self):
        return "<Sample('%s','%s')>" % (self.id, self.md5)

    def to_dict(self):
        """Converts object to dict.
        @return: dict
        """
        d = {}
        for column in self.__table__.columns:
            value = getattr(self, column.name)
            d[column.name] = value
        return d

    def to_json(self):
        """Converts object to JSON.
        @return: JSON data
        """
        return json.dumps(self.to_dict())

    def __init__(self,
                 md5,
                 crc32,
                 sha1,
                 sha256,
                 sha512,
                 file_size,
                 file_type=None,
                 ssdeep=None):
        self.md5 = md5
        self.sha1 = sha1
        self.crc32 = crc32
        self.sha256 = sha256
        self.sha512 = sha512
        self.file_size = file_size
        if file_type:
            self.file_type = file_type
        if ssdeep:
            self.ssdeep = ssdeep

class Database:
    """Analysis queue database.

    This class handles the creation of the database user for internal queue
    management. It also provides some functions for interacting with it.
    """

    def __init__(self, dsn=None):
        """@param dsn: database connection string."""
        cfg = Config()

        if dsn:
            self.engine = create_engine(dsn, poolclass=NullPool)
        elif cfg.cuckoo.database:
            self.engine = create_engine(cfg.cuckoo.database, poolclass=NullPool)
        else:
            db_file = os.path.join(CUCKOO_ROOT, "db", "cuckoo.db")
            if not os.path.exists(db_file):
                db_dir = os.path.dirname(db_file)
                if not os.path.exists(db_dir):
                    try:
                        create_folder(folder=db_dir)
                    except CuckooOperationalError as e:
                        raise CuckooDatabaseError("Unable to create database "
                                                  "directory: %s" % e)
            self.engine = create_engine("sqlite:///%s" % db_file)

        # Disable SQL logging. Turn it on for debugging.
        self.engine.echo = False
        # Connection timeout.
        if cfg.cuckoo.database_timeout:
            self.engine.pool_timeout = cfg.cuckoo.database_timeout
        else:
            self.engine.pool_timeout = 60
        # Create schema.
        try:
            Base.metadata.create_all(self.engine)
        except SQLAlchemyError as e:
            raise CuckooDatabaseError("Unable to create or connect to "
                                      "database: %s" % e)

        # Get db session.
        self.Session = sessionmaker(bind=self.engine)

    def __del__(self):
        """Disconnects pool."""
        self.engine.dispose()

    def _set_status(self, task_id, status):
        """Set task status.
        @param task_id: task identifier
        @param status: status string
        @return: operation status
        """
        session = self.Session()
        session.query(Task).get(task_id).status = status
        try:
            session.commit()
        except:
            session.rollback()
            return False

        return True

    def fetch(self):
        """Fetch a task.
        @return: task dict or None.
        """
        session = self.Session()
        row = session.query(Task).filter(Task.status == "pending").order_by("priority desc, added_on").first()
        return row

    def process(self, task_id):
        """Set task status as processing.
        @param task_id: task identifier
        @return: operation status
        """
        return self._set_status(task_id, "processing")

    def fetch_and_process(self):
        """Fetches a task waiting to be processed and locks it for processing.
        @return: None or task
        """
        session = self.Session()
        try:
            row = session.query(Task).filter(Task.status == "pending").order_by("priority desc, added_on").first()
            if row:
               row.status = "processing"
            session.commit()
        except SQLAlchemyError:
            session.rollback()
            return None
        return row

    def complete(self, task_id, success=True):
        """Mark a task as completed.
        @param task_id: task id.
        @param success: completed with status.
        @return: operation status.
        """
        session = self.Session()
        task = session.query(Task).get(task_id)

        if success:
            task.status = "success"
        else:
            task.status = "failure"

        task.completed_on = datetime.now()

        try:
            session.commit()
        except:
            session.rollback()
            return False

        return True

    def guest_start(self, task_id, name, label, manager):
        """Logs guest start.
        @param task_id: task identifier
        @param name: vm name
        @param label: vm label
        @param manager: vm manager
        @return: guest row id
        """
        session = self.Session()
        guest = Guest(name, label, manager)
        guest.started_on = datetime.now()
        session.query(Task).get(task_id).guest = guest
        try:
            session.commit()
        except:
            session.rollback()
            return None

        return guest.id

    def guest_stop(self, guest_id):
        """Logs guest stop.
        @param guest_id: guest log entry id
        """
        session = self.Session()
        session.query(Guest).get(guest_id).shutdown_on = datetime.now()
        try:
            session.commit()
        except:
            session.rollback()

    # The following functions are mostly used by external utils.

    def add(self,
            obj,
            timeout=0,
            package=None,
            options=None,
            priority=1,
            custom=None,
            machine=None,
            platform=None):
        """Add a task to database.
        @param file_path: sample path.
        @param timeout: selected timeout.
        @param options: analysis options.
        @param priority: analysis priority.
        @param custom: custom options.
        @param machine: selected machine.
        @param platform: platform
        @return: cursor or None.
        """
        session = self.Session()

        if isinstance(obj, File):
            try:
                sample = Sample(md5=obj.get_md5(),
                                crc32=obj.get_crc32(),
                                sha1=obj.get_sha1(),
                                sha256=obj.get_sha256(),
                                sha512=obj.get_sha512(),
                                file_size=obj.get_size(),
                                file_type=obj.get_type(),
                                ssdeep=obj.get_ssdeep())
                session.add(sample)
                session.commit()
            except IntegrityError:
                session.rollback()
                sample = session.query(Sample).filter(Sample.md5 == obj.get_md5()).first()

            task = Task(obj.file_path)
            task.sample_id = sample.id
        elif isinstance(obj, URL):
            task = Task(obj.url)

        task.category = obj.__class__.__name__.lower()
        task.timeout = timeout
        task.package = package
        task.options = options
        task.priority = priority
        task.custom = custom
        task.machine = machine
        task.platform = platform
        session.add(task)

        try:
            session.commit()
        except:
            session.rollback()
            return None

        return task.id

    def add_path(self,
                 file_path,
                 timeout=0,
                 package=None,
                 options=None,
                 priority=1,
                 custom=None,
                 machine=None,
                 platform=None):
        """Add a task to database from file path.
        @param file_path: sample path.
        @param timeout: selected timeout.
        @param options: analysis options.
        @param priority: analysis priority.
        @param custom: custom options.
        @param machine: selected machine.
        @param platform: platform
        @return: cursor or None.
        """
        if not file_path or not os.path.exists(file_path):
            return None

        return self.add(File(file_path),
                        timeout,
                        package,
                        options,
                        priority,
                        custom,
                        machine,
                        platform)

    def add_url(self,
                url,
                timeout=0,
                package=None,
                options=None,
                priority=1,
                custom=None,
                machine=None,
                platform=None):
        """Add a task to database from url.
        @param url: url.
        @param timeout: selected timeout.
        @param options: analysis options.
        @param priority: analysis priority.
        @param custom: custom options.
        @param machine: selected machine.
        @param platform: platform
        @return: cursor or None.
        """
        return self.add(URL(url),
                        timeout,
                        package,
                        options,
                        priority,
                        custom,
                        machine,
                        platform)

    def list_tasks(self, limit=None):
        """Retrieve list of task.
        @param limit: specify a limit of entries.
        @return: list of tasks.
        """
        session = self.Session()
        tasks = session.query(Task).order_by("added_on desc").limit(limit)
        return tasks

    def view_task(self, task_id):
        """Retrieve information on a task.
        @param task_id: ID of the task to query.
        @return: details on the task.
        """
        session = self.Session()
        task = session.query(Task).get(task_id)
        return task

    def view_sample(self, sample_id):
        """Retrieve information on a sample.
        @param sample_id: ID of the sample to query.
        @return: details on the sample.
        """
        session = self.Session()
        sample = session.query(Sample).get(sample_id)
        return sample

    def find_sample_by_md5(self, md5):
        """Search samples by MD5.
        @param md5: md5 string
        @return: matches list
        """
        session = self.Session()
        samples = session.query(Sample).filter(Sample.md5 == md5)
        return samples

    def find_sample_by_id(self, sample_id):
        """Search samples by ID.
        @param sample_id: id
        @return: matching object
        """
        session = self.Session()
        sample = session.query(Sample).get(sample_id)
        return sample
<commit_msg>Removed duplicated method<commit_after># Copyright (C) 2010-2012 Cuckoo Sandbox Developers.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import sys
import json
from datetime import datetime

from lib.cuckoo.common.constants import CUCKOO_ROOT
from lib.cuckoo.common.exceptions import CuckooDatabaseError
from lib.cuckoo.common.exceptions import CuckooOperationalError
from lib.cuckoo.common.exceptions import CuckooDependencyError
from lib.cuckoo.common.config import Config
from lib.cuckoo.common.objects import File, URL
from lib.cuckoo.common.utils import create_folder

try:
    from sqlalchemy import create_engine, Column
    from sqlalchemy import Integer, String, Boolean, DateTime, Enum
    from sqlalchemy import ForeignKey, Text, Index
    from sqlalchemy.orm import sessionmaker, relationship
    from sqlalchemy.sql import func
    from sqlalchemy.ext.declarative import declarative_base
    from sqlalchemy.exc import SQLAlchemyError, IntegrityError
    from sqlalchemy.pool import NullPool
    Base = declarative_base()
except ImportError:
    raise CuckooDependencyError("SQLAlchemy library not found, "
                                "verify your setup")

class Task(Base):
    """Analysis task queue."""
    __tablename__ = "tasks"

    id = Column(Integer(), primary_key=True)
    target = Column(Text(), nullable=False)
    category = Column(String(255), nullable=False)
    timeout = Column(Integer(), server_default="0", nullable=False)
    priority = Column(Integer(), server_default="1", nullable=False)
    custom = Column(String(255), nullable=True)
    machine = Column(String(255), nullable=True)
    package = Column(String(255), nullable=True)
    options = Column(String(255), nullable=True)
    platform = Column(String(255), nullable=True)
    added_on = Column(DateTime(timezone=False),
                      default=datetime.now(),
                      nullable=False)
    completed_on = Column(DateTime(timezone=False), nullable=True)
    status = Column(Enum("pending",
                         "processing",
                         "failure",
                         "success",
                         name="status_type"),
                         server_default="pending",
                         nullable=False)
    sample_id = Column(Integer, ForeignKey("samples.id"), nullable=True)
    sample = relationship("Sample", backref="tasks")
    guest = relationship("Guest", uselist=False, backref="tasks")

    def to_dict(self):
        """Converts object to dict.
        @return: dict
        """
        d = {}
        for column in self.__table__.columns:
            value = getattr(self, column.name)
            if isinstance(value, datetime):
                d[column.name] = value.strftime("%Y-%m-%d %H:%M:%S")
            else:
                d[column.name] = value
        return d

    def to_json(self):
        """Converts object to JSON.
        @return: JSON data
        """
        return json.dumps(self.to_dict())

    def __init__(self, target=None):
        self.target = target

    def __repr__(self):
        return "<Task('%s','%s')>" % (self.id, self.target)

class Guest(Base):
    """Tracks guest run."""
    __tablename__ = "guests"

    id = Column(Integer(), primary_key=True)
    name = Column(String(255), nullable=False)
    label = Column(String(255), nullable=False)
    manager = Column(String(255), nullable=False)
    started_on = Column(DateTime(timezone=False),
                        default=datetime.now(),
                        nullable=False)
    shutdown_on = Column(DateTime(timezone=False), nullable=True)
    task_id = Column(Integer,
                     ForeignKey('tasks.id'),
                     nullable=False,
                     unique=True)

    def __repr__(self):
        return "<Guest('%s','%s')>" % (self.id, self.name)

    def to_dict(self):
        """Converts object to dict.
        @return: dict
        """
        d = {}
        for column in self.__table__.columns:
            value = getattr(self, column.name)
            if isinstance(value, datetime):
                d[column.name] = value.strftime("%Y-%m-%d %H:%M:%S")
            else:
                d[column.name] = value

    def to_json(self):
        """Converts object to JSON.
        @return: JSON data
        """
        return json.dumps(self.to_dict())

    def __init__(self, name, label, manager):
        self.name = name
        self.label = label
        self.manager = manager

class Sample(Base):
    """Submitted files details."""
    __tablename__ = "samples"

    id = Column(Integer(), primary_key=True)
    file_size = Column(Integer(), nullable=False)
    file_type = Column(String(255), nullable=False)
    md5 = Column(String(32), nullable=False)
    crc32 = Column(String(8), nullable=False)
    sha1 = Column(String(40), nullable=False)
    sha256 = Column(String(64), nullable=False)
    sha512 = Column(String(128), nullable=False)
    ssdeep = Column(String(255), nullable=True)
    __table_args__ = (Index("hash_index",
                            "md5",
                            "crc32",
                            "sha1",
                            "sha256",
                            "sha512",
                            unique=True), )

    def __repr__(self):
        return "<Sample('%s','%s')>" % (self.id, self.md5)

    def to_dict(self):
        """Converts object to dict.
        @return: dict
        """
        d = {}
        for column in self.__table__.columns:
            value = getattr(self, column.name)
            d[column.name] = value
        return d

    def to_json(self):
        """Converts object to JSON.
        @return: JSON data
        """
        return json.dumps(self.to_dict())

    def __init__(self,
                 md5,
                 crc32,
                 sha1,
                 sha256,
                 sha512,
                 file_size,
                 file_type=None,
                 ssdeep=None):
        self.md5 = md5
        self.sha1 = sha1
        self.crc32 = crc32
        self.sha256 = sha256
        self.sha512 = sha512
        self.file_size = file_size
        if file_type:
            self.file_type = file_type
        if ssdeep:
            self.ssdeep = ssdeep

class Database:
    """Analysis queue database.

    This class handles the creation of the database user for internal queue
    management. It also provides some functions for interacting with it.
    """

    def __init__(self, dsn=None):
        """@param dsn: database connection string."""
        cfg = Config()

        if dsn:
            self.engine = create_engine(dsn, poolclass=NullPool)
        elif cfg.cuckoo.database:
            self.engine = create_engine(cfg.cuckoo.database, poolclass=NullPool)
        else:
            db_file = os.path.join(CUCKOO_ROOT, "db", "cuckoo.db")
            if not os.path.exists(db_file):
                db_dir = os.path.dirname(db_file)
                if not os.path.exists(db_dir):
                    try:
                        create_folder(folder=db_dir)
                    except CuckooOperationalError as e:
                        raise CuckooDatabaseError("Unable to create database "
                                                  "directory: %s" % e)
            self.engine = create_engine("sqlite:///%s" % db_file)

        # Disable SQL logging. Turn it on for debugging.
        self.engine.echo = False
        # Connection timeout.
        if cfg.cuckoo.database_timeout:
            self.engine.pool_timeout = cfg.cuckoo.database_timeout
        else:
            self.engine.pool_timeout = 60
        # Create schema.
        try:
            Base.metadata.create_all(self.engine)
        except SQLAlchemyError as e:
            raise CuckooDatabaseError("Unable to create or connect to "
                                      "database: %s" % e)

        # Get db session.
        self.Session = sessionmaker(bind=self.engine)

    def __del__(self):
        """Disconnects pool."""
        self.engine.dispose()

    def _set_status(self, task_id, status):
        """Set task status.
        @param task_id: task identifier
        @param status: status string
        @return: operation status
        """
        session = self.Session()
        session.query(Task).get(task_id).status = status
        try:
            session.commit()
        except:
            session.rollback()
            return False

        return True

    def fetch(self):
        """Fetch a task.
        @return: task dict or None.
        """
        session = self.Session()
        row = session.query(Task).filter(Task.status == "pending").order_by("priority desc, added_on").first()
        return row

    def process(self, task_id):
        """Set task status as processing.
        @param task_id: task identifier
        @return: operation status
        """
        return self._set_status(task_id, "processing")

    def fetch_and_process(self):
        """Fetches a task waiting to be processed and locks it for processing.
        @return: None or task
        """
        session = self.Session()
        try:
            row = session.query(Task).filter(Task.status == "pending").order_by("priority desc, added_on").first()
            if row:
               row.status = "processing"
            session.commit()
        except SQLAlchemyError:
            session.rollback()
            return None
        return row

    def complete(self, task_id, success=True):
        """Mark a task as completed.
        @param task_id: task id.
        @param success: completed with status.
        @return: operation status.
        """
        session = self.Session()
        task = session.query(Task).get(task_id)

        if success:
            task.status = "success"
        else:
            task.status = "failure"

        task.completed_on = datetime.now()

        try:
            session.commit()
        except:
            session.rollback()
            return False

        return True

    def guest_start(self, task_id, name, label, manager):
        """Logs guest start.
        @param task_id: task identifier
        @param name: vm name
        @param label: vm label
        @param manager: vm manager
        @return: guest row id
        """
        session = self.Session()
        guest = Guest(name, label, manager)
        guest.started_on = datetime.now()
        session.query(Task).get(task_id).guest = guest
        try:
            session.commit()
        except:
            session.rollback()
            return None

        return guest.id

    def guest_stop(self, guest_id):
        """Logs guest stop.
        @param guest_id: guest log entry id
        """
        session = self.Session()
        session.query(Guest).get(guest_id).shutdown_on = datetime.now()
        try:
            session.commit()
        except:
            session.rollback()

    # The following functions are mostly used by external utils.

    def add(self,
            obj,
            timeout=0,
            package=None,
            options=None,
            priority=1,
            custom=None,
            machine=None,
            platform=None):
        """Add a task to database.
        @param file_path: sample path.
        @param timeout: selected timeout.
        @param options: analysis options.
        @param priority: analysis priority.
        @param custom: custom options.
        @param machine: selected machine.
        @param platform: platform
        @return: cursor or None.
        """
        session = self.Session()

        if isinstance(obj, File):
            try:
                sample = Sample(md5=obj.get_md5(),
                                crc32=obj.get_crc32(),
                                sha1=obj.get_sha1(),
                                sha256=obj.get_sha256(),
                                sha512=obj.get_sha512(),
                                file_size=obj.get_size(),
                                file_type=obj.get_type(),
                                ssdeep=obj.get_ssdeep())
                session.add(sample)
                session.commit()
            except IntegrityError:
                session.rollback()
                sample = session.query(Sample).filter(Sample.md5 == obj.get_md5()).first()

            task = Task(obj.file_path)
            task.sample_id = sample.id
        elif isinstance(obj, URL):
            task = Task(obj.url)

        task.category = obj.__class__.__name__.lower()
        task.timeout = timeout
        task.package = package
        task.options = options
        task.priority = priority
        task.custom = custom
        task.machine = machine
        task.platform = platform
        session.add(task)

        try:
            session.commit()
        except:
            session.rollback()
            return None

        return task.id

    def add_path(self,
                 file_path,
                 timeout=0,
                 package=None,
                 options=None,
                 priority=1,
                 custom=None,
                 machine=None,
                 platform=None):
        """Add a task to database from file path.
        @param file_path: sample path.
        @param timeout: selected timeout.
        @param options: analysis options.
        @param priority: analysis priority.
        @param custom: custom options.
        @param machine: selected machine.
        @param platform: platform
        @return: cursor or None.
        """
        if not file_path or not os.path.exists(file_path):
            return None

        return self.add(File(file_path),
                        timeout,
                        package,
                        options,
                        priority,
                        custom,
                        machine,
                        platform)

    def add_url(self,
                url,
                timeout=0,
                package=None,
                options=None,
                priority=1,
                custom=None,
                machine=None,
                platform=None):
        """Add a task to database from url.
        @param url: url.
        @param timeout: selected timeout.
        @param options: analysis options.
        @param priority: analysis priority.
        @param custom: custom options.
        @param machine: selected machine.
        @param platform: platform
        @return: cursor or None.
        """
        return self.add(URL(url),
                        timeout,
                        package,
                        options,
                        priority,
                        custom,
                        machine,
                        platform)

    def list_tasks(self, limit=None):
        """Retrieve list of task.
        @param limit: specify a limit of entries.
        @return: list of tasks.
        """
        session = self.Session()
        tasks = session.query(Task).order_by("added_on desc").limit(limit)
        return tasks

    def view_task(self, task_id):
        """Retrieve information on a task.
        @param task_id: ID of the task to query.
        @return: details on the task.
        """
        session = self.Session()
        task = session.query(Task).get(task_id)
        return task

    def view_sample(self, sample_id):
        """Retrieve information on a sample.
        @param sample_id: ID of the sample to query.
        @return: details on the sample.
        """
        session = self.Session()
        sample = session.query(Sample).get(sample_id)
        return sample

    def find_sample_by_md5(self, md5):
        """Search samples by MD5.
        @param md5: md5 string
        @return: matches list
        """
        session = self.Session()
        samples = session.query(Sample).filter(Sample.md5 == md5)
        return samples
<|endoftext|>


<commit_before>#!/usr/bin/python
"""Lightbox library for JTAG's RGBController

This module contains the JSON-RPC web-interface for Lightbox
"""
__author__ = 'Elmer de Looff <elmer@underdark.nl>'
__version__ = '2.0'

# Standard modules
import BaseHTTPServer
import datetime
import mimetypes
import os
import simplejson
import sys

# Package modules
from . import light
from . import utils

class ApiHandler(BaseHTTPServer.BaseHTTPRequestHandler):
  """Ligtbox JSON API Handler."""
  def do_GET(self):
    """Very basic request router."""
    path = self.path
    if path == '/':
      return self._Redirect('/static/api.html')
    elif path == '/api':
      return self.ControllerInfo()
    elif path == '/api/outputs':
      return self.OutputInfo()
    elif path.startswith('/static/'):
      return self.ServeStatic()
    return self._ErrorResponse('No path %r. Try the root please.' % path)

  def _ErrorResponse(self, error):
    """Something didn't quite go as planned, let's tell the client something."""
    self.send_response(400)
    self.send_header('content-type', 'text/plain')
    self.end_headers()
    self.wfile.write(error)

  def _JsonResponse(self, data):
    """Successful request, send response to client as JSON."""
    return self._SuccessResponse(simplejson.dumps(data), 'application/json')

  def _SuccessResponse(self, data, content_type):
    """Returns a 200 OK with the given data and content-type."""
    self.send_response(200)
    self.send_header('content-type', content_type)
    self.end_headers()
    self.wfile.write(data)

  def _Redirect(self, location, code=301):
    """Redirects the client to the new location."""
    self.send_response(code)
    self.send_header('location', location)
    self.end_headers()
    self.wfile.write(location)

  def ControllerInfo(self):
    """Returns a JSON object with controller information."""
    self._JsonResponse(self.server.box.Info())

  def ServeStatic(self):
    """Returns files from the 'static' directory."""
    requested = os.path.abspath(self.path)
    if not requested.startswith('/static'):
      return self._ErrorResponse('Request not permitted: %r' % self.path)
    static_path = os.path.join(os.path.dirname(__file__), requested[1:])
    try:
      with file(static_path) as static_file:
        content_type, _encoding = mimetypes.guess_type(static_path)
        if not content_type:
          content_type = 'text/plain'
        return self._SuccessResponse(static_file.read(), content_type)
    except IOError:
      self.log_error('File not found: %r', self.path)
      return self._ErrorResponse('File not found: %r' % self.path)

  def OutputInfo(self):
    """Returns a JSON object with Lightbox output information."""
    outputs = []
    for output_id, output in enumerate(self.server.box):
      outputs.append({
          'outputNumber': output_id,
          'mixedColorRgb': output.color,
          'mixedColorHex': '#%02x%02x%02x' % tuple(output.color),
          'layers': list(LayerReport(output))})
    self._JsonResponse(outputs)

  def do_POST(self):
    """Processes Lightbox controls via JSON."""
    if self.path != '/api':
      self.log_error('Received POST on address other than /api.')
      return self._ErrorResponse('Can only write to /api.')
    content_type = self.headers['content-type'].split(';')[0]
    if content_type != 'application/json':
      self.log_error(
          'Received POST with unsupported content-type %r.', content_type)
      return self._ErrorResponse('Only application/json is allowed for POST.')
    if 'content-length' not in self.headers:
      self.log_error('Received POST without a specified content-length.')
      return self._ErrorResponse('Headers must provide the message length.')
    payload = simplejson.loads(
        self.rfile.read(int(self.headers['content-length'])))
    if isinstance(payload, list):
      map(self.ProcessCommand, payload)
    else:
      self.ProcessCommand(payload)
    self.send_response(200)
    self.end_headers()

  def ProcessCommand(self, api_command):
    """Performs the given command on the Lightbox instance.

    Envelope and blend method are loaded by name from the utils module. The
    handling method is selected from a string as well, defaulting to 'Fade' if
    none is provided.
    """
    command = api_command.copy()
    if 'blender' in command:
      if not hasattr(utils.Blenders, command['blender']):
        raise ValueError('Provided blender %r is not a known blender.' % (
            command['blender']))
      command['blender'] = getattr(utils.Blenders, command['blender'])
    if 'envelope' in command:
      if not hasattr(utils.Envelopes, command['envelope']):
        raise ValueError('Provided envelope %r is not a known envelope.' % (
            command['envelope']))
      command['envelope'] = getattr(utils.Envelopes, command['envelope'])
    channel = self.server.box[command.get('output', 0)]
    api_command['action'] = action = command.get('action', 'fade').capitalize()
    if not hasattr(light.ActionsMixIn, action):
      raise ValueError(
          'Chosen action %r is not an action for this channel.' %  action)
    getattr(channel, action)(**command)

  def log_error(self, fmt, *args):
    """Logs an error by prefixing 'Error' and sending it to log_message."""
    self.log_message('Error: %s' % fmt, *args)

  def log_message(self, fmt, *args):
    """Logs the messages from the RequestHandler

    This logs the requesting host (without reverse dns lookup), the time
    as ISO-8601 time string and the provided format and args.
    """
    if self.server.verbose:
      sys.stderr.write("%s - [%s] %s\n" % (
          self.client_address[0],
          datetime.datetime.now().strftime('%F %T.%f'),
          fmt % args))

def ApiServer(box, port=8000, quiet=False):
  """Starts and runs a JSON API server for the given Lightbox controller."""
  server = BaseHTTPServer.HTTPServer(('0.0.0.0', port), ApiHandler)
  server.box = box
  server.verbose = not quiet
  server.serve_forever()

def LayerReport(output):
  """Yields a dictionary with the state of each layer in the given output."""
  for layer in output:
    yield {'blender': layer.blender.__name__,
           'envelope': layer.envelope.__name__,
           'colorRgb': layer.color,
           'colorHex': '#%02x%02x%02x' % layer.color,
           'opacity': layer.opacity}
<commit_msg>Set max-age on static content returns.<commit_after>#!/usr/bin/python
"""Lightbox library for JTAG's RGBController

This module contains the JSON-RPC web-interface for Lightbox
"""
__author__ = 'Elmer de Looff <elmer@underdark.nl>'
__version__ = '2.0'

# Standard modules
import BaseHTTPServer
import datetime
import mimetypes
import os
import simplejson
import sys

# Package modules
from . import light
from . import utils

class ApiHandler(BaseHTTPServer.BaseHTTPRequestHandler):
  """Ligtbox JSON API Handler."""
  def do_GET(self):
    """Very basic request router."""
    path = self.path
    if path == '/':
      return self._Redirect('/static/api.html')
    elif path == '/api':
      return self.ControllerInfo()
    elif path == '/api/outputs':
      return self.OutputInfo()
    elif path.startswith('/static/'):
      return self.ServeStatic()
    return self._ErrorResponse('No path %r. Try the root please.' % path)

  def _ErrorResponse(self, error):
    """Something didn't quite go as planned, let's tell the client something."""
    self.send_response(400)
    self.send_header('content-type', 'text/plain')
    self.end_headers()
    self.wfile.write(error)

  def _JsonResponse(self, data):
    """Successful request, send response to client as JSON."""
    return self._SuccessResponse(simplejson.dumps(data), 'application/json')

  def _SuccessResponse(self, data, content_type, max_age=0):
    """Returns a 200 OK with the given data and content-type."""
    self.send_response(200)
    self.send_header('content-type', content_type)
    self.send_header('max-age', max_age)
    self.end_headers()
    self.wfile.write(data)

  def _Redirect(self, location, code=301):
    """Redirects the client to the new location."""
    self.send_response(code)
    self.send_header('location', location)
    self.end_headers()
    self.wfile.write(location)

  def ControllerInfo(self):
    """Returns a JSON object with controller information."""
    self._JsonResponse(self.server.box.Info())

  def ServeStatic(self):
    """Returns files from the 'static' directory."""
    requested = os.path.abspath(self.path)
    if not requested.startswith('/static'):
      return self._ErrorResponse('Request not permitted: %r' % self.path)
    static_path = os.path.join(os.path.dirname(__file__), requested[1:])
    try:
      with file(static_path) as static_file:
        content_type, _encoding = mimetypes.guess_type(static_path)
        if not content_type:
          content_type = 'text/plain'
        return self._SuccessResponse(
            static_file.read(), content_type, max_age=86400)
    except IOError:
      self.log_error('File not found: %r', self.path)
      return self._ErrorResponse('File not found: %r' % self.path)

  def OutputInfo(self):
    """Returns a JSON object with Lightbox output information."""
    outputs = []
    for output_id, output in enumerate(self.server.box):
      outputs.append({
          'outputNumber': output_id,
          'mixedColorRgb': output.color,
          'mixedColorHex': '#%02x%02x%02x' % tuple(output.color),
          'layers': list(LayerReport(output))})
    self._JsonResponse(outputs)

  def do_POST(self):
    """Processes Lightbox controls via JSON."""
    if self.path != '/api':
      self.log_error('Received POST on address other than /api.')
      return self._ErrorResponse('Can only write to /api.')
    content_type = self.headers['content-type'].split(';')[0]
    if content_type != 'application/json':
      self.log_error(
          'Received POST with unsupported content-type %r.', content_type)
      return self._ErrorResponse('Only application/json is allowed for POST.')
    if 'content-length' not in self.headers:
      self.log_error('Received POST without a specified content-length.')
      return self._ErrorResponse('Headers must provide the message length.')
    payload = simplejson.loads(
        self.rfile.read(int(self.headers['content-length'])))
    if isinstance(payload, list):
      map(self.ProcessCommand, payload)
    else:
      self.ProcessCommand(payload)
    self.send_response(200)
    self.end_headers()

  def ProcessCommand(self, api_command):
    """Performs the given command on the Lightbox instance.

    Envelope and blend method are loaded by name from the utils module. The
    handling method is selected from a string as well, defaulting to 'Fade' if
    none is provided.
    """
    command = api_command.copy()
    if 'blender' in command:
      if not hasattr(utils.Blenders, command['blender']):
        raise ValueError('Provided blender %r is not a known blender.' % (
            command['blender']))
      command['blender'] = getattr(utils.Blenders, command['blender'])
    if 'envelope' in command:
      if not hasattr(utils.Envelopes, command['envelope']):
        raise ValueError('Provided envelope %r is not a known envelope.' % (
            command['envelope']))
      command['envelope'] = getattr(utils.Envelopes, command['envelope'])
    channel = self.server.box[command.get('output', 0)]
    api_command['action'] = action = command.get('action', 'fade').capitalize()
    if not hasattr(light.ActionsMixIn, action):
      raise ValueError(
          'Chosen action %r is not an action for this channel.' %  action)
    getattr(channel, action)(**command)

  def log_error(self, fmt, *args):
    """Logs an error by prefixing 'Error' and sending it to log_message."""
    self.log_message('Error: %s' % fmt, *args)

  def log_message(self, fmt, *args):
    """Logs the messages from the RequestHandler

    This logs the requesting host (without reverse dns lookup), the time
    as ISO-8601 time string and the provided format and args.
    """
    if self.server.verbose:
      sys.stderr.write("%s - [%s] %s\n" % (
          self.client_address[0],
          datetime.datetime.now().strftime('%F %T.%f'),
          fmt % args))

def ApiServer(box, port=8000, quiet=False):
  """Starts and runs a JSON API server for the given Lightbox controller."""
  server = BaseHTTPServer.HTTPServer(('0.0.0.0', port), ApiHandler)
  server.box = box
  server.verbose = not quiet
  server.serve_forever()

def LayerReport(output):
  """Yields a dictionary with the state of each layer in the given output."""
  for layer in output:
    yield {'blender': layer.blender.__name__,
           'envelope': layer.envelope.__name__,
           'colorRgb': layer.color,
           'colorHex': '#%02x%02x%02x' % layer.color,
           'opacity': layer.opacity}
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import sys
import time
import os

import cv2
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import patches
from matplotlib.widgets import Slider, Button
import numpy as np
from scipy.misc import imread, imsave
from scipy.signal import medfilt, argrelmin
from skimage import filters, measure, color
from skimage.transform import iradon, downscale_local_mean

class TomoScan(object):

    def __init__(self, num_projections, save_folder, camera_port=0,
                 wait=0, save=False):
        """
        Note that for the subsequent functionality to work, the number of
        projections must be great enough to ensure that a rotational
        range > 360 degrees is captured.
        """
        self.folder = save_folder
        self.p0 = 0
        self.cor_offset = 0
        self.num_images = 0
        self.angles = None

        camera = cv2.VideoCapture(camera_port)
        retval, im = camera.read()
        try:
            dims = im[:, :, 2].shape + (num_projections, )
        except TypeError:
            error = (r'Camera returning None. Check camera settings (port) and'
                     r' ensure camera is not being run by other software.')
            raise TypeError(error)
        self.im_stack = np.zeros(dims)

        for i in range(num_projections):
            retval, im = camera.read()
            self.im_stack[:, :, i] = color.rgb2hsv(im)[:, :, 2]
            sys.stdout.write("\rProgress: [{0:20s}] {1:.0f}%".format('#' *
                             int(20*(i + 1) / num_projections),
                             100*((i + 1)/num_projections)))
            sys.stdout.flush()
            time.sleep(wait)
        del camera

        self.height = self.im_stack.shape[0]
        self.width = self.im_stack.shape[1]
        self.cor_offset = 0
        self.cropped = self.im_stack
        self.recon_data = None

        if save:
            save_folder = os.path.join(self.folder, 'projections')
            if not os.path.exists(save_folder):
                os.makedirs(save_folder)
            for idx in range(self.im_stack.shape[-1]):
                fpath = os.path.join(save_folder, '%04d.tif' % idx)
                imsave(fpath, self.im_stack[:, :, idx])

    def plot_histogram(self, proj=5):
        """
        Plots histogram of pixel intensity for specified projection.

        # proj:       Projection number (int)
        """
        histogram = np.histogram(self.im_stack[:, :, proj], 255)
        plt.plot(histogram[0])

    def auto_set_angles(self, proj_ref=5, order=25, plot=True):
        """
        Attempts to automatically locate image at 360 degrees (and multiples
        of 360 degrees). Alignment based on difference calculation between
        reference projection each subsequent projections.

        # proj_ref:   Projection to use as initial or reference projection.
                      Recommended to be greater than 1 (due to acquisition
                      spacing issues in initial projections)
        # order:      Window in which to search for minimas in the
                      difference calculations - should be approx equal to
                      (number of projections in 360) / 2
        # plot:       Plot the difference results
        """
        ref = downscale_local_mean(self.im_stack[:, :, proj_ref], (3, 3))
        diff = np.nan * np.ones((self.im_stack.shape[-1] - proj_ref))
        proj_nums = range(proj_ref, self.im_stack.shape[-1])
        for idx, i in enumerate(proj_nums):
            current = downscale_local_mean(self.im_stack[:, :, i], (3, 3))
            tmp = current - ref
            diff[idx] = tmp.std()
        minimas = argrelmin(np.array(diff), order=order)
        print(minimas)
        self.p0 = proj_ref
        self.num_images = minimas[0][0] + 1
        self.angles = np.linspace(0, 360, self.num_images, dtype=int)

        if plot:
            plt.figure()
            plt.plot(proj_nums, diff)
            plt.plot(minimas[0] + proj_ref, np.array(diff)[minimas], 'r*')
            plt.plot([minimas[0][0] + proj_ref, minimas[0][0] + proj_ref],
                     [0, np.max(diff)], 'r--')
            plt.xlabel('Image number')
            plt.ylabel('Thresholded Pixels Relative to Image 1')
            plt.text(minimas[0][0] + proj_ref, np.max(diff), r'$360^{\circ}$',
                     horizontalalignment='center', verticalalignment='bottom')

        print('\n%i images in a 360 rotation. \n\n If this is incorrect '
              'either rerun with a different value for order or use the manual'
              ' method.' % self.num_images)

    def auto_centre(self, window=400, cropped=False):
        """
        Automatic method for finding the centre of rotation.

        # window:     Window width to search across (pixels).
        """
        half_win = window // 2
        win_range = range(-half_win, half_win)

        if cropped:
            ref = self.cropped[:, half_win:-half_win, self.p0]
            im_180 = self.cropped[:, :, int(self.num_images / 2) + self.p0]
        else:
            ref = self.im_stack[:, half_win:-half_win, self.p0]
            im_180 = self.im_stack[:, :, int(self.num_images / 2) + self.p0]
        flipped = np.fliplr(im_180)

        diff = np.nan * np.zeros(len(win_range))

        for idx, i in enumerate(win_range):

            cropped = flipped[:, half_win + i: -half_win + i]
            tmp = cropped - ref
            diff[idx] = tmp.std()

        minima = np.argmin(diff)
        self.cor_offset = win_range[minima]

        plt.plot(win_range, diff)
        plt.plot(self.cor_offset, np.min(diff), '*')
        plt.ylabel('Standard deviation (original v 180deg flipped)')
        plt.xlabel('Cropped pixels')

        fig, ax_array = plt.subplots(1, 2, figsize=(10, 6))
        image = np.copy(self.im_stack[:, :, self.p0])
        if self.cor_offset <= 0:
            poly_pnts = [[self.width + self.cor_offset, 0], [self.width, 0],
                         [self.width, self.height],
                         [self.width + self.cor_offset, self.height]]
        else:
            poly_pnts = [[0, 0], [self.cor_offset, 0],
                         [self.cor_offset, self.height], [0, self.height]]
        ax_array[0].imshow(image)
        centre = self.width / 2 - self.cor_offset / 2
        ax_array[0].plot([centre, centre], [0, self.height], 'k-',
                         linewidth=2, label='New COR')
        ax_array[0].plot([self.width / 2, self.width / 2],
                         [0, self.height], 'r-', linewidth=2, label='Old COR')
        ax_array[0].legend()
        ax_array[0].set_xlim([0, image.shape[1]])
        ax_array[0].set_ylim([image.shape[0], 0])

        if self.cor_offset <= 0:
            image = np.copy(self.im_stack[:, -self.cor_offset:, self.p0])
        else:
            image = np.copy(self.im_stack[:, :-self.cor_offset, self.p0])

        ax_array[1].imshow(image)
        ax_array[1].plot([image.shape[1]/2, image.shape[1]/2],
                         [0, self.height], 'k-', linewidth=2, label='New COR')
        ax_array[1].legend()

        ax_array[1].set_xlim([0, image.shape[1]])
        ax_array[1].set_ylim([image.shape[0], 0])

        ax_array[0].add_patch(patches.Polygon(poly_pnts, closed=True,
                              fill=False, hatch='///', color='k'))
        ax_array[0].set_title('Uncropped')
        ax_array[1].set_title('Cropped and centred')

    def crop_to_centre(self):
        """
        Temporary method for re-cropping of data after finding the cor_offset.
        """
        if self.cor_offset < 0:
            self.cropped = self.im_stack[:, -self.cor_offset:, :]
        else:
            self.cropped = self.im_stack[:, :-self.cor_offset, :]

    def manual_set_angles(self, interact=True, proj_ref=5,
                          num_images=None, ang_range=None):
        """
        Manually define the number of images in 360 degrees. Defaults to
        interactive mode in which images can be compared against initial,
        reference image.

        # interact:   Run in interactive mode (True/False)
        # proj_ref:   Projection to use as initial or reference projection.
                      Recommended to be greater than 1 (due to acquisiton
                      spacing issues in intital projections)
        # num_images: If not in interact mode, manually specify number
                      of images
        # ang_range:  If not in interact mode, manually specify angular range
                      of images (must be multiple of 180)
        """
        if interact:
            backend = matplotlib.get_backend()
            err = ("Matplotlib running inline. Plot interaction not possible."
                   "\nTry running %matplotlib in the ipython console (and "
                   "%matplotlib inline to return to default behaviour). In "
                   "standard console use matplotlib.use('TkAgg') to interact.")

            assert backend != 'module://ipykernel.pylab.backend_inline', err
            fig, ax_array = plt.subplots(1, 2, figsize=(10, 5))

            ax_slider = plt.axes([0.2, 0.07, 0.5, 0.05])
            ax_button = plt.axes([0.81, 0.05, 0.1, 0.075])

            ax_array[0].imshow(self.im_stack[:, :, proj_ref])
            ax_array[1].imshow(self.im_stack[:, :, proj_ref])
            ax_array[0].axis('off')
            ax_array[1].axis('off')
            fig.tight_layout()
            fig.subplots_adjust(bottom=0.2)
            nfiles = self.im_stack.shape[-1] + 1
            window_slider = Slider(ax_slider, 'Image', proj_ref,
                                   nfiles, valinit=0)
            store_button = Button(ax_button, r'Save - 360')

            def slider_update(val):
                ax_array[1].imshow(self.im_stack[:, :, int(window_slider.val)])
                window_slider.valtext.set_text('%i' % window_slider.val)
                fig.canvas.draw_idle()

            window_slider.on_changed(slider_update)

            def store_data(label):
                # Check this is correct - proj_ref!!!
                self.num_images = int(window_slider.val) - proj_ref + 1
                self.angles = np.linspace(0, 360, self.num_images)
                self.p0 = proj_ref
                plt.close()

            store_button.on_clicked(store_data)
            return window_slider, store_button

        else:
            error = 'Images must cover a rotational range of 180 or 360 deg'
            assert (ang_range == 180) or (ang_range == 360), error
            self.angles = np.linspace(0, ang_range, num_images)
            self.im_stack = self.im_stack[:, :, :num_images]

    def reconstruct(self, downsample=(4, 4, 1), pre_filter=True,
                    kernel=9, save=True):

        if self.cor_offset <= 0:
            images = self.im_stack[:, -self.cor_offset:,
                                   self.p0:self.num_images + self.p0]
        else:
            images = self.im_stack[:, :-self.cor_offset,
                                   self.p0:self.num_images + self.p0]

        images = downscale_local_mean(images, downsample)

        if pre_filter is not False:
            for i in range(images.shape[-1]):
                images[:, :, i] = medfilt(images[:, :, i], kernel_size=kernel)

        for j in range(self.height):
            sinotmp = np.squeeze(images[j, :, :])
            imagetmp = iradon(sinotmp, theta=self.angles,
                              filter=None, circle=True)

            self.recon_data[:, :, j] = imagetmp
            if save:
                save_folder = os.path.join(self.folder, 'reconstruction')
                if not os.path.exists(save_folder):
                    os.makedirs(save_folder)
                fpath = os.path.join(save_folder, '%04d.tif' % j)
                imsave(fpath, imagetmp)

class LoadScan(TomoScan):

    def __init__(self, folder):
        self.folder = folder
        files = [f for f in os.listdir(folder) if f[-4:] == '.tif']
        im_shape = imread(os.path.join(self.folder, files[0])).shape
        self.im_stack = np.zeros(im_shape + (len(files), ))
        for idx, fname in enumerate(files):
            sys.stdout.write("\rProgress: [{0:20s}] {1:.0f}%".format('#' *
                             int(20*(idx + 1) / len(files)),
                             100*((idx + 1)/len(files))))
            sys.stdout.flush()
            f = os.path.join(self.folder, fname)
            self.im_stack[:, :, idx] = imread(f)
        self.cor_offset = 0
<commit_msg>Tidying code. Removed defaults and unused code.<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import sys
import time
import os

import cv2
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import patches
from matplotlib.widgets import Slider, Button
import numpy as np
from scipy.misc import imread, imsave
from scipy.signal import medfilt, argrelmin
from skimage import color
from skimage.transform import iradon, downscale_local_mean

class TomoScan(object):

    def __init__(self, num_projections, save_folder, camera_port=0,
                 wait=0, save=False):
        """
        Note that for the subsequent functionality to work, the number of
        projections must be great enough to ensure that a rotational
        range > 360 degrees is captured.
        """
        self.folder = save_folder
        self.p0 = 0
        self.cor_offset = 0
        self.num_images = 0
        self.angles = None
        self.recon_data = None

        camera = cv2.VideoCapture(camera_port)
        retval, im = camera.read()
        try:
            dims = im[:, :, 2].shape + (num_projections, )
        except TypeError:
            error = (r'Camera returning None. Check camera settings (port) and'
                     r' ensure camera is not being run by other software.')
            raise TypeError(error)
        self.im_stack = np.zeros(dims)

        for i in range(num_projections):
            retval, im = camera.read()
            self.im_stack[:, :, i] = color.rgb2hsv(im)[:, :, 2]
            sys.stdout.write("\rProgress: [{0:20s}] {1:.0f}%".format('#' *
                             int(20*(i + 1) / num_projections),
                             100*((i + 1)/num_projections)))
            sys.stdout.flush()
            time.sleep(wait)
        del camera

        self.height = self.im_stack.shape[0]
        self.width = self.im_stack.shape[1]

        if save:
            save_folder = os.path.join(self.folder, 'projections')
            if not os.path.exists(save_folder):
                os.makedirs(save_folder)
            for idx in range(self.im_stack.shape[-1]):
                fpath = os.path.join(save_folder, '%04d.tif' % idx)
                imsave(fpath, self.im_stack[:, :, idx])

    def plot_histogram(self, proj=5):
        """
        Plots histogram of pixel intensity for specified projection.

        # proj:       Projection number (int)
        """
        histogram = np.histogram(self.im_stack[:, :, proj], 255)
        plt.plot(histogram[0])

    def auto_set_angles(self, order, proj_ref=5, plot=True):
        """
        Attempts to automatically locate image at 360 degrees (and multiples
        of 360 degrees). Alignment based on difference calculation between
        reference projection each subsequent projections.

        # proj_ref:   Projection to use as initial or reference projection.
                      Recommended to be greater than 1 (due to acquisition
                      spacing issues in initial projections)
        # order:      Window in which to search for minimas in the
                      difference calculations - should be approx equal to
                      (number of projections in 360) / 2
        # plot:       Plot the difference results
        """
        ref = downscale_local_mean(self.im_stack[:, :, proj_ref], (3, 3))
        diff = np.nan * np.ones((self.im_stack.shape[-1] - proj_ref))
        proj_nums = range(proj_ref, self.im_stack.shape[-1])
        for idx, i in enumerate(proj_nums):
            current = downscale_local_mean(self.im_stack[:, :, i], (3, 3))
            tmp = current - ref
            diff[idx] = tmp.std()
        minimas = argrelmin(np.array(diff), order=order)
        print(minimas)
        self.p0 = proj_ref
        self.num_images = minimas[0][0] + 1
        self.angles = np.linspace(0, 360, self.num_images, dtype=int)

        if plot:
            plt.figure()
            plt.plot(proj_nums, diff)
            plt.plot(minimas[0] + proj_ref, np.array(diff)[minimas], 'r*')
            plt.plot([minimas[0][0] + proj_ref, minimas[0][0] + proj_ref],
                     [0, np.max(diff)], 'r--')
            plt.xlabel('Image number')
            plt.ylabel('Thresholded Pixels Relative to Image 1')
            plt.text(minimas[0][0] + proj_ref, np.max(diff), r'$360^{\circ}$',
                     horizontalalignment='center', verticalalignment='bottom')

        print('\n%i images in a 360 rotation. \n\n If this is incorrect '
              'either rerun with a different value for order or use the manual'
              ' method.' % self.num_images)

    def auto_centre(self, window=400):
        """
        Automatic method for finding the centre of rotation.

        # window:     Window width to search across (pixels).
        """
        half_win = window // 2
        win_range = range(-half_win, half_win)

        ref = self.im_stack[:, half_win:-half_win, self.p0]
        im_180 = self.im_stack[:, :, int(self.num_images / 2) + self.p0]
        flipped = np.fliplr(im_180)

        diff = np.nan * np.zeros(len(win_range))

        for idx, i in enumerate(win_range):

            cropped = flipped[:, half_win + i: -half_win + i]
            tmp = cropped - ref
            diff[idx] = tmp.std()

        minima = np.argmin(diff)
        self.cor_offset = win_range[minima]

        plt.plot(win_range, diff)
        plt.plot(self.cor_offset, np.min(diff), '*')
        plt.ylabel('Standard deviation (original v 180deg flipped)')
        plt.xlabel('Cropped pixels')

        fig, ax_array = plt.subplots(1, 2, figsize=(10, 6))
        image = np.copy(self.im_stack[:, :, self.p0])
        if self.cor_offset <= 0:
            poly_pnts = [[self.width + self.cor_offset, 0], [self.width, 0],
                         [self.width, self.height],
                         [self.width + self.cor_offset, self.height]]
        else:
            poly_pnts = [[0, 0], [self.cor_offset, 0],
                         [self.cor_offset, self.height], [0, self.height]]
        ax_array[0].imshow(image)
        centre = self.width / 2 - self.cor_offset / 2
        ax_array[0].plot([centre, centre], [0, self.height], 'k-',
                         linewidth=2, label='New COR')
        ax_array[0].plot([self.width / 2, self.width / 2],
                         [0, self.height], 'r-', linewidth=2, label='Old COR')
        ax_array[0].legend()
        ax_array[0].set_xlim([0, image.shape[1]])
        ax_array[0].set_ylim([image.shape[0], 0])

        if self.cor_offset <= 0:
            image = np.copy(self.im_stack[:, -self.cor_offset:, self.p0])
        else:
            image = np.copy(self.im_stack[:, :-self.cor_offset, self.p0])

        ax_array[1].imshow(image)
        ax_array[1].plot([image.shape[1]/2, image.shape[1]/2],
                         [0, self.height], 'k-', linewidth=2, label='New COR')
        ax_array[1].legend()

        ax_array[1].set_xlim([0, image.shape[1]])
        ax_array[1].set_ylim([image.shape[0], 0])

        ax_array[0].add_patch(patches.Polygon(poly_pnts, closed=True,
                              fill=False, hatch='///', color='k'))
        ax_array[0].set_title('Uncropped')
        ax_array[1].set_title('Cropped and centred')

    def manual_set_angles(self, interact=True, proj_ref=5,
                          num_images=None, ang_range=None):
        """
        Manually define the number of images in 360 degrees. Defaults to
        interactive mode in which images can be compared against initial,
        reference image.

        # interact:   Run in interactive mode (True/False)
        # proj_ref:   Projection to use as initial or reference projection.
                      Recommended to be greater than 1 (due to acquisiton
                      spacing issues in intital projections)
        # num_images: If not in interact mode, manually specify number
                      of images
        # ang_range:  If not in interact mode, manually specify angular range
                      of images (must be multiple of 180)
        """
        if interact:
            backend = matplotlib.get_backend()
            err = ("Matplotlib running inline. Plot interaction not possible."
                   "\nTry running %matplotlib in the ipython console (and "
                   "%matplotlib inline to return to default behaviour). In "
                   "standard console use matplotlib.use('TkAgg') to interact.")

            assert backend != 'module://ipykernel.pylab.backend_inline', err
            fig, ax_array = plt.subplots(1, 2, figsize=(10, 5))

            ax_slider = plt.axes([0.2, 0.07, 0.5, 0.05])
            ax_button = plt.axes([0.81, 0.05, 0.1, 0.075])

            ax_array[0].imshow(self.im_stack[:, :, proj_ref])
            ax_array[1].imshow(self.im_stack[:, :, proj_ref])
            ax_array[0].axis('off')
            ax_array[1].axis('off')
            fig.tight_layout()
            fig.subplots_adjust(bottom=0.2)
            nfiles = self.im_stack.shape[-1] + 1
            window_slider = Slider(ax_slider, 'Image', proj_ref,
                                   nfiles, valinit=0)
            store_button = Button(ax_button, r'Save - 360')

            def slider_update(val):
                ax_array[1].imshow(self.im_stack[:, :, int(window_slider.val)])
                window_slider.valtext.set_text('%i' % window_slider.val)
                fig.canvas.draw_idle()

            window_slider.on_changed(slider_update)

            def store_data(label):
                # Check this is correct - proj_ref!!!
                self.num_images = int(window_slider.val) - proj_ref + 1
                self.angles = np.linspace(0, 360, self.num_images)
                self.p0 = proj_ref
                plt.close()

            store_button.on_clicked(store_data)
            return window_slider, store_button

        else:
            error = 'Images must cover a rotational range of 180 or 360 deg'
            assert (ang_range == 180) or (ang_range == 360), error
            self.angles = np.linspace(0, ang_range, num_images)
            self.im_stack = self.im_stack[:, :, :num_images]

    def reconstruct(self, downsample=(4, 4, 1), pre_filter=True,
                    kernel=9, save=True):

        if self.cor_offset <= 0:
            images = self.im_stack[:, -self.cor_offset:,
                                   self.p0:self.num_images + self.p0]
        else:
            images = self.im_stack[:, :-self.cor_offset,
                                   self.p0:self.num_images + self.p0]

        images = downscale_local_mean(images, downsample)

        if pre_filter is not False:
            for i in range(images.shape[-1]):
                images[:, :, i] = medfilt(images[:, :, i], kernel_size=kernel)

        for j in range(self.height):
            sinotmp = np.squeeze(images[j, :, :])
            imagetmp = iradon(sinotmp, theta=self.angles,
                              filter=None, circle=True)

            self.recon_data[:, :, j] = imagetmp
            if save:
                save_folder = os.path.join(self.folder, 'reconstruction')
                if not os.path.exists(save_folder):
                    os.makedirs(save_folder)
                fpath = os.path.join(save_folder, '%04d.tif' % j)
                imsave(fpath, imagetmp)

class LoadProjections(TomoScan):

    def __init__(self, load_folder):
        self.folder = load_folder
        files = [f for f in os.listdir(load_folder) if f[-4:] == '.tif']
        im_shape = imread(os.path.join(self.folder, files[0])).shape
        self.im_stack = np.zeros(im_shape + (len(files), ))
        for idx, fname in enumerate(files):
            sys.stdout.write("\rProgress: [{0:20s}] {1:.0f}%".format('#' *
                             int(20*(idx + 1) / len(files)),
                             100*((idx + 1)/len(files))))
            sys.stdout.flush()
            f = os.path.join(self.folder, fname)
            self.im_stack[:, :, idx] = imread(f)

        self.p0 = 0
        self.cor_offset = 0
        self.num_images = 0
        self.angles = None
        self.height = self.im_stack.shape[0]
        self.width = self.im_stack.shape[1]
<|endoftext|>


<commit_before>from django.db import models

class Todo(models.Model):
    name = models.CharField(max_length=200)
    due_date = models.DateTimeField('due date')
    completed_date = models.DateTimeField('completed date')

    def __unicode__(self):
        return self.name
<commit_msg>Use python3 syntax<commit_after>from django.db import models

class Todo(models.Model):
    name = models.CharField(max_length=200)
    due_date = models.DateTimeField('due date')
    completed_date = models.DateTimeField('completed date')

    def __str__(self):
        return self.name
<|endoftext|>


<commit_before># Do not edit, pipeline versioning governed by git tags!
__version__ = '1.4.12-12-g10b3fe4d'
<commit_msg>[#1.4.13] updated version number<commit_after># Do not edit, pipeline versioning governed by git tags!
__version__ = '1.4.13'
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

#  Licensed under the Apache License, Version 2.0 (the "License"); you may
#  not use this file except in compliance with the License. You may obtain
#  a copy of the License at
#
#       https://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#  License for the specific language governing permissions and limitations
#  under the License.

"""Meta data of line-bot-sdk."""

__version__ = '2.0.0'
__author__ = 'LINE Corporation'
__copyright__ = 'Copyright 2016, LINE Corporation'
__license__ = 'Apache 2.0'

__all__ = (
    '__version__'
)
<commit_msg>Release 2.0.1 (#345)<commit_after># -*- coding: utf-8 -*-

#  Licensed under the Apache License, Version 2.0 (the "License"); you may
#  not use this file except in compliance with the License. You may obtain
#  a copy of the License at
#
#       https://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#  License for the specific language governing permissions and limitations
#  under the License.

"""Meta data of line-bot-sdk."""

__version__ = '2.0.1'
__author__ = 'LINE Corporation'
__copyright__ = 'Copyright 2016, LINE Corporation'
__license__ = 'Apache 2.0'

__all__ = (
    '__version__'
)
<|endoftext|>