<commit_before>import json
import os
from collections import namedtuple
import zipfile

from .utils import extract_archive, verify_str_arg, iterable_to_str
from .vision import VisionDataset
from PIL import Image

class Cityscapes(VisionDataset):
    """`Cityscapes <http://www.cityscapes-dataset.com/>`_ Dataset.

    Args:
        root (string): Root directory of dataset where directory ``leftImg8bit``
            and ``gtFine`` or ``gtCoarse`` are located.
        split (string, optional): The image split to use, ``train``, ``test`` or ``val`` if mode="fine"
            otherwise ``train``, ``train_extra`` or ``val``
        mode (string, optional): The quality mode to use, ``fine`` or ``coarse``
        target_type (string or list, optional): Type of target to use, ``instance``, ``semantic``, ``polygon``
            or ``color``. Can also be a list to output a tuple with all specified target types.
        transform (callable, optional): A function/transform that takes in a PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
        transforms (callable, optional): A function/transform that takes input sample and its target as entry
            and returns a transformed version.

    Examples:

        Get semantic segmentation target

        .. code-block:: python

            dataset = Cityscapes('./data/cityscapes', split='train', mode='fine',
                                 target_type='semantic')

            img, smnt = dataset[0]

        Get multiple targets

        .. code-block:: python

            dataset = Cityscapes('./data/cityscapes', split='train', mode='fine',
                                 target_type=['instance', 'color', 'polygon'])

            img, (inst, col, poly) = dataset[0]

        Validate on the "coarse" set

        .. code-block:: python

            dataset = Cityscapes('./data/cityscapes', split='val', mode='coarse',
                                 target_type='semantic')

            img, smnt = dataset[0]
    """

    # Based on https://github.com/mcordts/cityscapesScripts
    CityscapesClass = namedtuple('CityscapesClass', ['name', 'id', 'train_id', 'category', 'category_id',
                                                     'has_instances', 'ignore_in_eval', 'color'])

    classes = [
        CityscapesClass('unlabeled', 0, 255, 'void', 0, False, True, (0, 0, 0)),
        CityscapesClass('ego vehicle', 1, 255, 'void', 0, False, True, (0, 0, 0)),
        CityscapesClass('rectification border', 2, 255, 'void', 0, False, True, (0, 0, 0)),
        CityscapesClass('out of roi', 3, 255, 'void', 0, False, True, (0, 0, 0)),
        CityscapesClass('static', 4, 255, 'void', 0, False, True, (0, 0, 0)),
        CityscapesClass('dynamic', 5, 255, 'void', 0, False, True, (111, 74, 0)),
        CityscapesClass('ground', 6, 255, 'void', 0, False, True, (81, 0, 81)),
        CityscapesClass('road', 7, 0, 'flat', 1, False, False, (128, 64, 128)),
        CityscapesClass('sidewalk', 8, 1, 'flat', 1, False, False, (244, 35, 232)),
        CityscapesClass('parking', 9, 255, 'flat', 1, False, True, (250, 170, 160)),
        CityscapesClass('rail track', 10, 255, 'flat', 1, False, True, (230, 150, 140)),
        CityscapesClass('building', 11, 2, 'construction', 2, False, False, (70, 70, 70)),
        CityscapesClass('wall', 12, 3, 'construction', 2, False, False, (102, 102, 156)),
        CityscapesClass('fence', 13, 4, 'construction', 2, False, False, (190, 153, 153)),
        CityscapesClass('guard rail', 14, 255, 'construction', 2, False, True, (180, 165, 180)),
        CityscapesClass('bridge', 15, 255, 'construction', 2, False, True, (150, 100, 100)),
        CityscapesClass('tunnel', 16, 255, 'construction', 2, False, True, (150, 120, 90)),
        CityscapesClass('pole', 17, 5, 'object', 3, False, False, (153, 153, 153)),
        CityscapesClass('polegroup', 18, 255, 'object', 3, False, True, (153, 153, 153)),
        CityscapesClass('traffic light', 19, 6, 'object', 3, False, False, (250, 170, 30)),
        CityscapesClass('traffic sign', 20, 7, 'object', 3, False, False, (220, 220, 0)),
        CityscapesClass('vegetation', 21, 8, 'nature', 4, False, False, (107, 142, 35)),
        CityscapesClass('terrain', 22, 9, 'nature', 4, False, False, (152, 251, 152)),
        CityscapesClass('sky', 23, 10, 'sky', 5, False, False, (70, 130, 180)),
        CityscapesClass('person', 24, 11, 'human', 6, True, False, (220, 20, 60)),
        CityscapesClass('rider', 25, 12, 'human', 6, True, False, (255, 0, 0)),
        CityscapesClass('car', 26, 13, 'vehicle', 7, True, False, (0, 0, 142)),
        CityscapesClass('truck', 27, 14, 'vehicle', 7, True, False, (0, 0, 70)),
        CityscapesClass('bus', 28, 15, 'vehicle', 7, True, False, (0, 60, 100)),
        CityscapesClass('caravan', 29, 255, 'vehicle', 7, True, True, (0, 0, 90)),
        CityscapesClass('trailer', 30, 255, 'vehicle', 7, True, True, (0, 0, 110)),
        CityscapesClass('train', 31, 16, 'vehicle', 7, True, False, (0, 80, 100)),
        CityscapesClass('motorcycle', 32, 17, 'vehicle', 7, True, False, (0, 0, 230)),
        CityscapesClass('bicycle', 33, 18, 'vehicle', 7, True, False, (119, 11, 32)),
        CityscapesClass('license plate', -1, -1, 'vehicle', 7, False, True, (0, 0, 142)),
    ]

    def __init__(self, root, split='train', mode='fine', target_type='instance',
                 transform=None, target_transform=None, transforms=None):
        super(Cityscapes, self).__init__(root, transforms, transform, target_transform)
        self.mode = 'gtFine' if mode == 'fine' else 'gtCoarse'
        self.images_dir = os.path.join(self.root, 'leftImg8bit', split)
        self.targets_dir = os.path.join(self.root, self.mode, split)
        self.target_type = target_type
        self.split = split
        self.images = []
        self.targets = []

        verify_str_arg(mode, "mode", ("fine", "coarse"))
        if mode == "fine":
            valid_modes = ("train", "test", "val")
        else:
            valid_modes = ("train", "train_extra", "val")
        msg = ("Unknown value '{}' for argument split if mode is '{}'. "
               "Valid values are {{{}}}.")
        msg = msg.format(split, mode, iterable_to_str(valid_modes))
        verify_str_arg(split, "split", valid_modes, msg)

        if not isinstance(target_type, list):
            self.target_type = [target_type]
        [verify_str_arg(value, "target_type",
                        ("instance", "semantic", "polygon", "color"))
         for value in self.target_type]

        if not os.path.isdir(self.images_dir) or not os.path.isdir(self.targets_dir):

            if split == 'train_extra':
                image_dir_zip = os.path.join(self.root, 'leftImg8bit{}'.format('_trainextra.zip'))
            else:
                image_dir_zip = os.path.join(self.root, 'leftImg8bit{}'.format('_trainvaltest.zip'))

            if self.mode == 'gtFine':
                target_dir_zip = os.path.join(self.root, '{}{}'.format(self.mode, '_trainvaltest.zip'))
            elif self.mode == 'gtCoarse':
                target_dir_zip = os.path.join(self.root, '{}{}'.format(self.mode, '.zip'))

            if os.path.isfile(image_dir_zip) and os.path.isfile(target_dir_zip):
                extract_archive(from_path=image_dir_zip, to_path=self.root)
                extract_archive(from_path=target_dir_zip, to_path=self.root)
            else:
                raise RuntimeError('Dataset not found or incomplete. Please make sure all required folders for the'
                                   ' specified "split" and "mode" are inside the "root" directory')

        for city in os.listdir(self.images_dir):
            img_dir = os.path.join(self.images_dir, city)
            target_dir = os.path.join(self.targets_dir, city)
            for file_name in os.listdir(img_dir):
                target_types = []
                for t in self.target_type:
                    target_name = '{}_{}'.format(file_name.split('_leftImg8bit')[0],
                                                 self._get_target_suffix(self.mode, t))
                    target_types.append(os.path.join(target_dir, target_name))

                self.images.append(os.path.join(img_dir, file_name))
                self.targets.append(target_types)

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            tuple: (image, target) where target is a tuple of all target types if target_type is a list with more
            than one item. Otherwise target is a json object if target_type="polygon", else the image segmentation.
        """

        image = Image.open(self.images[index]).convert('RGB')

        targets = []
        for i, t in enumerate(self.target_type):
            if t == 'polygon':
                target = self._load_json(self.targets[index][i])
            else:
                target = Image.open(self.targets[index][i])

            targets.append(target)

        target = tuple(targets) if len(targets) > 1 else targets[0]

        if self.transforms is not None:
            image, target = self.transforms(image, target)

        return image, target

    def __len__(self):
        return len(self.images)

    def extra_repr(self):
        lines = ["Split: {split}", "Mode: {mode}", "Type: {target_type}"]
        return '\n'.join(lines).format(**self.__dict__)

    def _load_json(self, path):
        with open(path, 'r') as file:
            data = json.load(file)
        return data

    def _get_target_suffix(self, mode, target_type):
        if target_type == 'instance':
            return '{}_instanceIds.png'.format(mode)
        elif target_type == 'semantic':
            return '{}_labelIds.png'.format(mode)
        elif target_type == 'color':
            return '{}_color.png'.format(mode)
        else:
            return '{}_polygons.json'.format(mode)
<commit_msg>cityscapes (#2525)<commit_after>import json
import os
from collections import namedtuple
import zipfile
from typing import Any, Callable, Dict, List, Optional, Union, Tuple

from .utils import extract_archive, verify_str_arg, iterable_to_str
from .vision import VisionDataset
from PIL import Image

class Cityscapes(VisionDataset):
    """`Cityscapes <http://www.cityscapes-dataset.com/>`_ Dataset.

    Args:
        root (string): Root directory of dataset where directory ``leftImg8bit``
            and ``gtFine`` or ``gtCoarse`` are located.
        split (string, optional): The image split to use, ``train``, ``test`` or ``val`` if mode="fine"
            otherwise ``train``, ``train_extra`` or ``val``
        mode (string, optional): The quality mode to use, ``fine`` or ``coarse``
        target_type (string or list, optional): Type of target to use, ``instance``, ``semantic``, ``polygon``
            or ``color``. Can also be a list to output a tuple with all specified target types.
        transform (callable, optional): A function/transform that takes in a PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
        transforms (callable, optional): A function/transform that takes input sample and its target as entry
            and returns a transformed version.

    Examples:

        Get semantic segmentation target

        .. code-block:: python

            dataset = Cityscapes('./data/cityscapes', split='train', mode='fine',
                                 target_type='semantic')

            img, smnt = dataset[0]

        Get multiple targets

        .. code-block:: python

            dataset = Cityscapes('./data/cityscapes', split='train', mode='fine',
                                 target_type=['instance', 'color', 'polygon'])

            img, (inst, col, poly) = dataset[0]

        Validate on the "coarse" set

        .. code-block:: python

            dataset = Cityscapes('./data/cityscapes', split='val', mode='coarse',
                                 target_type='semantic')

            img, smnt = dataset[0]
    """

    # Based on https://github.com/mcordts/cityscapesScripts
    CityscapesClass = namedtuple('CityscapesClass', ['name', 'id', 'train_id', 'category', 'category_id',
                                                     'has_instances', 'ignore_in_eval', 'color'])

    classes = [
        CityscapesClass('unlabeled', 0, 255, 'void', 0, False, True, (0, 0, 0)),
        CityscapesClass('ego vehicle', 1, 255, 'void', 0, False, True, (0, 0, 0)),
        CityscapesClass('rectification border', 2, 255, 'void', 0, False, True, (0, 0, 0)),
        CityscapesClass('out of roi', 3, 255, 'void', 0, False, True, (0, 0, 0)),
        CityscapesClass('static', 4, 255, 'void', 0, False, True, (0, 0, 0)),
        CityscapesClass('dynamic', 5, 255, 'void', 0, False, True, (111, 74, 0)),
        CityscapesClass('ground', 6, 255, 'void', 0, False, True, (81, 0, 81)),
        CityscapesClass('road', 7, 0, 'flat', 1, False, False, (128, 64, 128)),
        CityscapesClass('sidewalk', 8, 1, 'flat', 1, False, False, (244, 35, 232)),
        CityscapesClass('parking', 9, 255, 'flat', 1, False, True, (250, 170, 160)),
        CityscapesClass('rail track', 10, 255, 'flat', 1, False, True, (230, 150, 140)),
        CityscapesClass('building', 11, 2, 'construction', 2, False, False, (70, 70, 70)),
        CityscapesClass('wall', 12, 3, 'construction', 2, False, False, (102, 102, 156)),
        CityscapesClass('fence', 13, 4, 'construction', 2, False, False, (190, 153, 153)),
        CityscapesClass('guard rail', 14, 255, 'construction', 2, False, True, (180, 165, 180)),
        CityscapesClass('bridge', 15, 255, 'construction', 2, False, True, (150, 100, 100)),
        CityscapesClass('tunnel', 16, 255, 'construction', 2, False, True, (150, 120, 90)),
        CityscapesClass('pole', 17, 5, 'object', 3, False, False, (153, 153, 153)),
        CityscapesClass('polegroup', 18, 255, 'object', 3, False, True, (153, 153, 153)),
        CityscapesClass('traffic light', 19, 6, 'object', 3, False, False, (250, 170, 30)),
        CityscapesClass('traffic sign', 20, 7, 'object', 3, False, False, (220, 220, 0)),
        CityscapesClass('vegetation', 21, 8, 'nature', 4, False, False, (107, 142, 35)),
        CityscapesClass('terrain', 22, 9, 'nature', 4, False, False, (152, 251, 152)),
        CityscapesClass('sky', 23, 10, 'sky', 5, False, False, (70, 130, 180)),
        CityscapesClass('person', 24, 11, 'human', 6, True, False, (220, 20, 60)),
        CityscapesClass('rider', 25, 12, 'human', 6, True, False, (255, 0, 0)),
        CityscapesClass('car', 26, 13, 'vehicle', 7, True, False, (0, 0, 142)),
        CityscapesClass('truck', 27, 14, 'vehicle', 7, True, False, (0, 0, 70)),
        CityscapesClass('bus', 28, 15, 'vehicle', 7, True, False, (0, 60, 100)),
        CityscapesClass('caravan', 29, 255, 'vehicle', 7, True, True, (0, 0, 90)),
        CityscapesClass('trailer', 30, 255, 'vehicle', 7, True, True, (0, 0, 110)),
        CityscapesClass('train', 31, 16, 'vehicle', 7, True, False, (0, 80, 100)),
        CityscapesClass('motorcycle', 32, 17, 'vehicle', 7, True, False, (0, 0, 230)),
        CityscapesClass('bicycle', 33, 18, 'vehicle', 7, True, False, (119, 11, 32)),
        CityscapesClass('license plate', -1, -1, 'vehicle', 7, False, True, (0, 0, 142)),
    ]

    def __init__(
            self,
            root: str,
            split: str = "train",
            mode: str = "fine",
            target_type: Union[List[str], str] = "instance",
            transform: Optional[Callable] = None,
            target_transform: Optional[Callable] = None,
            transforms: Optional[Callable] = None,
    ) -> None:
        super(Cityscapes, self).__init__(root, transforms, transform, target_transform)
        self.mode = 'gtFine' if mode == 'fine' else 'gtCoarse'
        self.images_dir = os.path.join(self.root, 'leftImg8bit', split)
        self.targets_dir = os.path.join(self.root, self.mode, split)
        self.target_type = target_type
        self.split = split
        self.images = []
        self.targets = []

        verify_str_arg(mode, "mode", ("fine", "coarse"))
        if mode == "fine":
            valid_modes = ("train", "test", "val")
        else:
            valid_modes = ("train", "train_extra", "val")
        msg = ("Unknown value '{}' for argument split if mode is '{}'. "
               "Valid values are {{{}}}.")
        msg = msg.format(split, mode, iterable_to_str(valid_modes))
        verify_str_arg(split, "split", valid_modes, msg)

        if not isinstance(target_type, list):
            self.target_type = [target_type]
        [verify_str_arg(value, "target_type",
                        ("instance", "semantic", "polygon", "color"))
         for value in self.target_type]

        if not os.path.isdir(self.images_dir) or not os.path.isdir(self.targets_dir):

            if split == 'train_extra':
                image_dir_zip = os.path.join(self.root, 'leftImg8bit{}'.format('_trainextra.zip'))
            else:
                image_dir_zip = os.path.join(self.root, 'leftImg8bit{}'.format('_trainvaltest.zip'))

            if self.mode == 'gtFine':
                target_dir_zip = os.path.join(self.root, '{}{}'.format(self.mode, '_trainvaltest.zip'))
            elif self.mode == 'gtCoarse':
                target_dir_zip = os.path.join(self.root, '{}{}'.format(self.mode, '.zip'))

            if os.path.isfile(image_dir_zip) and os.path.isfile(target_dir_zip):
                extract_archive(from_path=image_dir_zip, to_path=self.root)
                extract_archive(from_path=target_dir_zip, to_path=self.root)
            else:
                raise RuntimeError('Dataset not found or incomplete. Please make sure all required folders for the'
                                   ' specified "split" and "mode" are inside the "root" directory')

        for city in os.listdir(self.images_dir):
            img_dir = os.path.join(self.images_dir, city)
            target_dir = os.path.join(self.targets_dir, city)
            for file_name in os.listdir(img_dir):
                target_types = []
                for t in self.target_type:
                    target_name = '{}_{}'.format(file_name.split('_leftImg8bit')[0],
                                                 self._get_target_suffix(self.mode, t))
                    target_types.append(os.path.join(target_dir, target_name))

                self.images.append(os.path.join(img_dir, file_name))
                self.targets.append(target_types)

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """
        Args:
            index (int): Index
        Returns:
            tuple: (image, target) where target is a tuple of all target types if target_type is a list with more
            than one item. Otherwise target is a json object if target_type="polygon", else the image segmentation.
        """

        image = Image.open(self.images[index]).convert('RGB')

        targets: Any = []
        for i, t in enumerate(self.target_type):
            if t == 'polygon':
                target = self._load_json(self.targets[index][i])
            else:
                target = Image.open(self.targets[index][i])

            targets.append(target)

        target = tuple(targets) if len(targets) > 1 else targets[0]

        if self.transforms is not None:
            image, target = self.transforms(image, target)

        return image, target

    def __len__(self) -> int:
        return len(self.images)

    def extra_repr(self) -> str:
        lines = ["Split: {split}", "Mode: {mode}", "Type: {target_type}"]
        return '\n'.join(lines).format(**self.__dict__)

    def _load_json(self, path: str) -> Dict[str, Any]:
        with open(path, 'r') as file:
            data = json.load(file)
        return data

    def _get_target_suffix(self, mode: str, target_type: str) -> str:
        if target_type == 'instance':
            return '{}_instanceIds.png'.format(mode)
        elif target_type == 'semantic':
            return '{}_labelIds.png'.format(mode)
        elif target_type == 'color':
            return '{}_color.png'.format(mode)
        else:
            return '{}_polygons.json'.format(mode)
<|endoftext|>


<commit_before>#!/usr/bin/env python3

# This requires a sqlite database dump from
# http://pozniak.pl/dbdump/ody101-sqlite3-v12.db.bz2

# Something that I ran to extract module lists from copy pasted forum posts:
# cat new-list | tr '[' '\n' | tr ',' '\n' | sed -e 's/^[ \t]*//' | sed 's/ x.$//' | sort | uniq | ./market-stuff.py --filter

import sys
import sqlite3
import urllib.request as urlreq
from xml.dom.minidom import parseString

CHUNK_SIZE = 100
SYSTEM = 'Hemin'
ITEM_LIST = 'items'

name2id = {}
id2name = {}

conn = sqlite3.connect('ody101-sqlite3-v12.db')

def get_system_id(name):
    c = conn.cursor()
    c.execute("SELECT itemID from invnames where itemName = ?", (name,))
    return c.fetchone()

system_id = get_system_id(SYSTEM)

def load_items():
    c = conn.cursor()
    c.execute("SELECT typeId, typeName from invTypes")

    for (itemid, name) in c:
#        print(itemid, name)
        name2id[name] = itemid
        id2name[itemid] = name

def download_data(ids):
    base_url = 'http://api.eve-central.com/api/marketstat?hours=24&usesystem=%d&' % system_id
    suffix = "&".join("typeid=%d" % i for i in ids)
    url = base_url + suffix
#    print(url)
    s = urlreq.urlopen(url)
    s = "".join(x.decode() for x in s)
#    print(s)

    obj = parseString(s)
    return obj

def chunk(l, size):
    chunked = []
    i = 0
    while i < len(l):
        chunked.append(l[i:i+size])
        i += size
    return chunked

def read_xml_field(node, key):
    return node.getElementsByTagName(key)[0].childNodes[0].data

def handle_data(table, xml):
    items = xml.getElementsByTagName("type")
    for item in items:
        i = int(item.getAttribute("id"))
        sell = item.getElementsByTagName("sell")[0]
        volume = int(read_xml_field(sell, "volume"))
        min_price = float(read_xml_field(sell, "min"))
        table.append((id2name[i], volume, min_price))

def text_output(table):
    for parts in table:
        print("%s: %d at %s" % parts)

def html_output(table):
    print("<title>%s market data</title>" % SYSTEM)
    print("""

<link rel="stylesheet" type="text/css" href="http://ajax.aspnetcdn.com/ajax/jquery.dataTables/1.9.4/css/jquery.dataTables.css">

""")
    print("Data may be out of date or missing. Items might be in the wrong station. Price shown is the lowest price. If you want more items on here, message sully on IRC with links to lists of items (probably pastebinned).
")

    print("
%s market
" % SYSTEM)

    print("
")
    print("

Item	quantity	price

")
    print("
")
    for (name, count, price) in table:
        price_fmt = "{:,.2f}".format(price)
        print("
%s	%d	%s
" % (name, count, price_fmt))
    print("

")

def make_table(formatter):
    item_names = [s.strip() for s in open(ITEM_LIST)]
#    print(item_names)
    item_ids = [name2id[name] for name in item_names]
#    print(item_ids)

    table = []
    for part in chunk(item_ids, CHUNK_SIZE):
        data = download_data(part)
        handle_data(table, data)

    formatter(table)

def make_poller():
    item_names = [s.strip() for s in open(ITEM_LIST)]
    item_ids = [name2id[name] for name in item_names]
    print("item_ids = %s;" % str(item_ids))

def filter_input():
    for name in sys.stdin:
        name = name.strip()
        if name in name2id:
            print(name)

def main(args):
    load_items()

    if len(args) > 1 and args[1] == "--filter":
        filter_input()
    elif len(args) > 1 and args[1] == "--poller":
        make_poller()
    elif len(args) > 1 and args[1] == "--text":
        make_table(text_output)
    else:
        make_table(html_output)

if __name__ == '__main__':
    sys.exit(main(sys.argv))
<commit_msg>fix sorting by price<commit_after>#!/usr/bin/env python3

# This requires a sqlite database dump from
# http://pozniak.pl/dbdump/ody101-sqlite3-v12.db.bz2

# Something that I ran to extract module lists from copy pasted forum posts:
# cat new-list | tr '[' '\n' | tr ',' '\n' | sed -e 's/^[ \t]*//' | sed 's/ x.$//' | sort | uniq | ./market-stuff.py --filter

import sys
import sqlite3
import urllib.request as urlreq
from xml.dom.minidom import parseString

CHUNK_SIZE = 100
SYSTEM = 'Hemin'
ITEM_LIST = 'items'

name2id = {}
id2name = {}

conn = sqlite3.connect('ody101-sqlite3-v12.db')

def get_system_id(name):
    c = conn.cursor()
    c.execute("SELECT itemID from invnames where itemName = ?", (name,))
    return c.fetchone()

system_id = get_system_id(SYSTEM)

def load_items():
    c = conn.cursor()
    c.execute("SELECT typeId, typeName from invTypes")

    for (itemid, name) in c:
#        print(itemid, name)
        name2id[name] = itemid
        id2name[itemid] = name

def download_data(ids):
    base_url = 'http://api.eve-central.com/api/marketstat?hours=24&usesystem=%d&' % system_id
    suffix = "&".join("typeid=%d" % i for i in ids)
    url = base_url + suffix
#    print(url)
    s = urlreq.urlopen(url)
    s = "".join(x.decode() for x in s)
#    print(s)

    obj = parseString(s)
    return obj

def chunk(l, size):
    chunked = []
    i = 0
    while i < len(l):
        chunked.append(l[i:i+size])
        i += size
    return chunked

def read_xml_field(node, key):
    return node.getElementsByTagName(key)[0].childNodes[0].data

def handle_data(table, xml):
    items = xml.getElementsByTagName("type")
    for item in items:
        i = int(item.getAttribute("id"))
        sell = item.getElementsByTagName("sell")[0]
        volume = int(read_xml_field(sell, "volume"))
        min_price = float(read_xml_field(sell, "min"))
        table.append((id2name[i], volume, min_price))

def text_output(table):
    for parts in table:
        print("%s: %d at %s" % parts)

def html_output(table):
    print("<title>%s market data</title>" % SYSTEM)
    print("""

<link rel="stylesheet" type="text/css" href="http://ajax.aspnetcdn.com/ajax/jquery.dataTables/1.9.4/css/jquery.dataTables.css">

""")
    print("Data may be out of date or missing. Items might be in the wrong station. Price shown is the lowest price. If you want more items on here, message sully on IRC with links to lists of items (probably pastebinned).
")

    print("
%s market
" % SYSTEM)

    print("
")
    print("

Item	quantity	price

")
    print("
")
    for (name, count, price) in table:
        price_fmt = "{:,.2f}".format(price)
        print("
%s	%d	%s
" % (name, count, price_fmt))
    print("

")

def make_table(formatter):
    item_names = [s.strip() for s in open(ITEM_LIST)]
#    print(item_names)
    item_ids = [name2id[name] for name in item_names]
#    print(item_ids)

    table = []
    for part in chunk(item_ids, CHUNK_SIZE):
        data = download_data(part)
        handle_data(table, data)

    formatter(table)

def make_poller():
    item_names = [s.strip() for s in open(ITEM_LIST)]
    item_ids = [name2id[name] for name in item_names]
    print("item_ids = %s;" % str(item_ids))

def filter_input():
    for name in sys.stdin:
        name = name.strip()
        if name in name2id:
            print(name)

def main(args):
    load_items()

    if len(args) > 1 and args[1] == "--filter":
        filter_input()
    elif len(args) > 1 and args[1] == "--poller":
        make_poller()
    elif len(args) > 1 and args[1] == "--text":
        make_table(text_output)
    else:
        make_table(html_output)

if __name__ == '__main__':
    sys.exit(main(sys.argv))
<|endoftext|>


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

#    Copyright 2013 Mirantis, Inc.
#
#    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
#
#         http://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.

import traceback

import web
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy.orm.query import Query
from sqlalchemy.exc import ProgrammingError
from sqlalchemy import create_engine

from nailgun.logger import logger
from nailgun.settings import settings

db_str = "{engine}://{user}:{passwd}@{host}:{port}/{name}".format(
    **settings.DATABASE)

def make_engine():
    return create_engine(db_str, client_encoding='utf8')

engine = make_engine()

class NoCacheQuery(Query):
    """
    Override for common Query class.
    Needed for automatic refreshing objects
    from database during every query for evading
    problems with multiple sessions
    """
    def __init__(self, *args, **kwargs):
        self._populate_existing = True
        super(NoCacheQuery, self).__init__(*args, **kwargs)

def make_session(custom_engine=None):
    session = scoped_session(
        sessionmaker(bind=(custom_engine or engine), query_cls=NoCacheQuery))
    return session

def orm():
    if not hasattr(web.ctx, "orm"):
        web.ctx.orm = make_session()

    return web.ctx.orm

def load_db_driver(handler):
    web.ctx.orm = make_session()
    try:
        return handler()
    except web.HTTPError:
        web.ctx.orm.commit()
        raise
    except:
        web.ctx.orm.rollback()
        raise
    finally:
        web.ctx.orm.commit()
        web.ctx.orm.expire_all()

def syncdb():
    from nailgun.api.models import Base
    Base.metadata.create_all(engine)

def dropdb():
    db = make_session()

    tables = [name for (name,) in db.execute(
        "SELECT tablename FROM pg_tables WHERE schemaname = 'public'")]
    for table in tables:
        db.execute("DROP TABLE IF EXISTS %s CASCADE" % table)

    # sql query to list all types, equivalent to psql's \dT+
    types = [name for (name,) in db.execute("""
        SELECT t.typname as type FROM pg_type t
        LEFT JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
        WHERE (t.typrelid = 0 OR (
            SELECT c.relkind = 'c' FROM pg_catalog.pg_class c
            WHERE c.oid = t.typrelid
        ))
        AND NOT EXISTS(
            SELECT 1 FROM pg_catalog.pg_type el
            WHERE el.oid = t.typelem AND el.typarray = t.oid
        )
        AND n.nspname = 'public'
        """)]
    for type_ in types:
        db.execute("DROP TYPE IF EXISTS %s CASCADE" % type_)
    db.commit()

def flush():
    import nailgun.api.models as models
    import sqlalchemy.ext.declarative as dec
    session = scoped_session(sessionmaker(bind=engine))
    for attr in dir(models):
        attr_impl = getattr(models, attr)
        if isinstance(attr_impl, dec.DeclarativeMeta) \
                and not attr_impl is models.Base:
            map(session.delete, session.query(attr_impl).all())
    # for table in reversed(models.Base.metadata.sorted_tables):
    #     session.execute(table.delete())
    session.commit()
<commit_msg>Create per thread sqlalchemy session and engine<commit_after># -*- coding: utf-8 -*-

#    Copyright 2013 Mirantis, Inc.
#
#    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
#
#         http://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.

import traceback

import web
import threading
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy.orm.query import Query
from sqlalchemy.exc import ProgrammingError
from sqlalchemy import create_engine

from nailgun.logger import logger
from nailgun.settings import settings

db_str = "{engine}://{user}:{passwd}@{host}:{port}/{name}".format(
    **settings.DATABASE)

thread_local_storage = threading.local()

def db():
    """
    Global db session per thread
    """
    if not hasattr(thread_local_storage, 'db'):
        thread_local_storage.db = make_session(get_engine())
    return thread_local_storage.db

def get_engine():
    """
    Global engine per thread
    """
    if not hasattr(thread_local_storage, 'db_engine'):
        thread_local_storage.db_engine = make_engine()

    return thread_local_storage.db_engine

def make_engine():
    return create_engine(db_str, client_encoding='utf8')

def make_session(custom_engine=None):
    session = scoped_session(
        sessionmaker(
            bind=(custom_engine or get_engine()), query_cls=NoCacheQuery))
    return session

class NoCacheQuery(Query):
    """
    Override for common Query class.
    Needed for automatic refreshing objects
    from database during every query for evading
    problems with multiple sessions
    """
    def __init__(self, *args, **kwargs):
        self._populate_existing = True
        super(NoCacheQuery, self).__init__(*args, **kwargs)

def load_db_driver(handler):
    try:
        return handler()
    except web.HTTPError:
        db().commit()
        raise
    except:
        db().rollback()
        raise
    finally:
        db().commit()
        db().expire_all()

def syncdb():
    from nailgun.api.models import Base
    Base.metadata.create_all(get_engine())

def dropdb():
    db = make_session()

    tables = [name for (name,) in db.execute(
        "SELECT tablename FROM pg_tables WHERE schemaname = 'public'")]
    for table in tables:
        db.execute("DROP TABLE IF EXISTS %s CASCADE" % table)

    # sql query to list all types, equivalent to psql's \dT+
    types = [name for (name,) in db.execute("""
        SELECT t.typname as type FROM pg_type t
        LEFT JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
        WHERE (t.typrelid = 0 OR (
            SELECT c.relkind = 'c' FROM pg_catalog.pg_class c
            WHERE c.oid = t.typrelid
        ))
        AND NOT EXISTS(
            SELECT 1 FROM pg_catalog.pg_type el
            WHERE el.oid = t.typelem AND el.typarray = t.oid
        )
        AND n.nspname = 'public'
        """)]
    for type_ in types:
        db.execute("DROP TYPE IF EXISTS %s CASCADE" % type_)
    db.commit()

def flush():
    import nailgun.api.models as models
    import sqlalchemy.ext.declarative as dec
    session = scoped_session(sessionmaker(bind=get_engine()))
    for attr in dir(models):
        attr_impl = getattr(models, attr)
        if isinstance(attr_impl, dec.DeclarativeMeta) \
                and not attr_impl is models.Base:
            map(session.delete, session.query(attr_impl).all())
    # for table in reversed(models.Base.metadata.sorted_tables):
    #     session.execute(table.delete())
    session.commit()
<|endoftext|>


<commit_before>#!/usr/bin/env python
import logging
import sys
sys.path.append( '.' )
import shutil
from patXML import SQLPatent
from patXML import XMLPatent
from patXML import uniasc
from fwork  import *
import os, datetime, re

flder = '/var/share/patentdata/patents/testem'
logfile = flder + "/" + 'xml-parsing.log'
logging.basicConfig(filename=logfile, level=logging.DEBUG)
#logging.debug('This message should go to the log file')
#logging.info('So should this')
#logging.warning('And this, too')

t1 = datetime.datetime.now()

#get a listing of all files within the directory that follow the naming pattern
files = [x for x in os.listdir(flder)
         #if re.match(r"ip[a-z]{2}[0-9]{6,8}[.]xml", x, re.I)!=None]
         #if re.match(r"ipg\d{6}.one.xml", x, re.I)!=None]
         if re.match(r"ipg\d{6}.xml", x, re.I)!=None]
print "Total files: %d" % (len(files))
logging.info("Total files: %d" % (len(files)))

#sys.exit()

tables = ["assignee", "citation", "class", "inventor", "patent", "patdesc", "lawyer", "sciref", "usreldoc"]
total_count = 0
total_patents = 0
for filenum, filename in enumerate(files):
    print " > Regular Expression: %s" % filename
    XMLs = re.findall(
            r"""
                ([<][?]xml[ ]version.*?[>]       #all XML starts with ?xml
                .*?
                [<][/]us[-]patent[-]grant[>])    #and here is the end tag
             """,
            open(flder+"/"+files[filenum]).read(), re.I + re.S + re.X)
    print "   - Total Patents: %d" % (len(XMLs))
    logging.info("   - Total Patents: %d" % (len(XMLs)))

    xmllist = []
    count = 0
    patents = 0

    for i, x in enumerate(XMLs):
        try:
            xmllist.append(XMLPatent(x))
            patents += 1
        except Exception as inst:
            print type(inst)
            logging.error(type(inst))
            print "  - Error: %s (%d)  %s" % (filename, i, x[175:200])
            logging.error("  - Error: %s (%d)  %s" % (filename, i, x[175:200]))
            count += 1

    print "   - number of patents:", len(xmllist), datetime.datetime.now()-t1
    logging.info("   - number of patents: %d %s ", len(xmllist), datetime.datetime.now()-t1)
    print "   - number of errors: ", count
    logging.info( "   - number of errors: %d", count)
    total_count += count
    total_patents += patents

    for table in tables:
        # Cut the chaining here to better parameterize the call, allowing
        # the databases to be built in place
        # (/var/share/patentdata/patents/<year>)
        # outdb = flder + "/" + table
        q = SQLPatent().tblBuild(xmllist, tbl=table)
        SQLPatent().dbBuild(q, tbl=table, week=filename)
        #SQLPatent().dbBuild(q=SQLPatent().tblBuild(xmllist, tbl=table), tbl=table, week=filename)

    print "   -", datetime.datetime.now()-t1
    logging.info("   - %s", datetime.datetime.now()-t1)
    logging.info("   - total errors: %d", total_count)
    logging.info("   - total patents: %d", total_patents)

for table in tables:
    SQLPatent().dbFinal(tbl=table)

for table in tables:
    filename = table + ".sqlite3"
    shutil.move(filename,flder)

<commit_msg>synching for merge<commit_after>#!/usr/bin/env python
import logging
import sys
sys.path.append( '.' )
import shutil
from patXML import SQLPatent
from patXML import XMLPatent
from patXML import uniasc
from fwork  import *
import os, datetime, re

flder = '/var/share/patentdata/patents/2009'
logfile = flder + "/" + 'xml-parsing.log'
logging.basicConfig(filename=logfile, level=logging.DEBUG)
#logging.debug('This message should go to the log file')
#logging.info('So should this')
#logging.warning('And this, too')

t1 = datetime.datetime.now()

#get a listing of all files within the directory that follow the naming pattern
files = [x for x in os.listdir(flder)
         #if re.match(r"ip[a-z]{2}[0-9]{6,8}[.]xml", x, re.I)!=None]
         #if re.match(r"ipg\d{6}.one.xml", x, re.I)!=None]
         if re.match(r"ipg\d{6}.xml", x, re.I)!=None]
print "Total files: %d" % (len(files))
logging.info("Total files: %d" % (len(files)))

#sys.exit()

tables = ["assignee", "citation", "class", "inventor", "patent", "patdesc", "lawyer", "sciref", "usreldoc"]
total_count = 0
total_patents = 0
for filenum, filename in enumerate(files):
    print " > Regular Expression: %s" % filename
    XMLs = re.findall(
            r"""
                ([<][?]xml[ ]version.*?[>]       #all XML starts with ?xml
                .*?
                [<][/]us[-]patent[-]grant[>])    #and here is the end tag
             """,
            open(flder+"/"+files[filenum]).read(), re.I + re.S + re.X)
    print "   - Total Patents: %d" % (len(XMLs))
    logging.info("   - Total Patents: %d" % (len(XMLs)))

    xmllist = []
    count = 0
    patents = 0

    for i, x in enumerate(XMLs):
        try:
            xmllist.append(XMLPatent(x))
            patents += 1
        except Exception as inst:
            print type(inst)
            logging.error(type(inst))
            print "  - Error: %s (%d)  %s" % (filename, i, x[175:200])
            logging.error("  - Error: %s (%d)  %s" % (filename, i, x[175:200]))
            count += 1

    print "   - number of patents:", len(xmllist), datetime.datetime.now()-t1
    logging.info("   - number of patents: %d %s ", len(xmllist), datetime.datetime.now()-t1)
    print "   - number of errors: ", count
    logging.info( "   - number of errors: %d", count)
    total_count += count
    total_patents += patents

    for table in tables:
        # Cut the chaining here to better parameterize the call, allowing
        # the databases to be built in place
        # (/var/share/patentdata/patents/<year>)
        # outdb = flder + "/" + table
        q = SQLPatent().tblBuild(xmllist, tbl=table)
        SQLPatent().dbBuild(q, tbl=table, week=filename)
        #SQLPatent().dbBuild(q=SQLPatent().tblBuild(xmllist, tbl=table), tbl=table, week=filename)

    print "   -", datetime.datetime.now()-t1
    logging.info("   - %s", datetime.datetime.now()-t1)
    logging.info("   - total errors: %d", total_count)
    logging.info("   - total patents: %d", total_patents)

for table in tables:
    SQLPatent().dbFinal(tbl=table)

for table in tables:
    filename = table + ".sqlite3"
    shutil.move(filename,flder)

<|endoftext|>


<commit_before>
import numpy as np
import astropy.units as u
from spectral_cube import SpectralCube
# from astropy.utils.console import ProgressBar
import sys
from warnings import warn

from bubble_segment2D import BubbleFinder2D
from bubble_objects import Bubble3D
from bubble_catalog import PPV_Catalog
from clustering import cluster_brute_force, threeD_overlaps
from utils import sig_clip
from galaxy_utils import gal_props_checker
from progressbar import ProgressBar

class BubbleFinder(object):
    """docstring for BubbleFinder"""
    def __init__(self, cube, wcs=None, mask=None, sigma=None, empty_channel=0,
                 keep_threshold_mask=True, distance=None, galaxy_props={}):
        super(BubbleFinder, self).__init__()

        if not isinstance(cube, SpectralCube):
            if wcs is None:
                raise TypeError("When cube is not a SpectralCube, wcs must be"
                                " given.")
            cube = SpectralCube(cube, wcs)
            if mask is not None:
                cube = cube.with_mask(mask)

        self.cube = cube

        self.empty_channel = empty_channel

        if sigma is None:
            self.estimate_sigma()
        else:
            self.sigma = sigma

        self.keep_threshold_mask = keep_threshold_mask
        self._mask = None
        self.distance = distance
        self.galaxy_props = galaxy_props

    @property
    def cube(self):
        return self._cube

    @cube.setter
    def cube(self, input_cube):
        if input_cube.ndim != 3:
            raise TypeError("Given cube must be 3D.")

        self._cube = input_cube

    @property
    def mask(self):
        if not self.keep_threshold_mask:
            warn("Enable keep_threshold_mask to access the entire mask.")

        return self._mask

    def estimate_sigma(self, nsig=10):
        '''
        Use empty channels to estimate sigma. Uses iterative sigma clipping
        to obtain a robust estimate.
        '''

        self.sigma = sig_clip(self.cube[self.empty_channel], nsig=nsig)

    @property
    def sigma(self):
        return self._sigma

    @sigma.setter
    def sigma(self, val):
        if val < 0.0 or ~np.isfinite(val):
            raise ValueError("sigma must be positive and finite.")

        self._sigma = val

    @property
    def galaxy_props(self):
        return self._galaxy_props

    @galaxy_props.setter
    def galaxy_props(self, input_dict):

        # Make sure all of the kwargs are given.
        gal_props_checker(input_dict)

        self._galaxy_props = input_dict

    def get_bubbles(self, verbose=True, overlap_frac=0.9, min_channels=3,
                    use_cube_mask=False, nsig=2., refit=False, scales=None,
                    cube_linewidth=None, multiprocess=True, nprocesses=None,
                    overlap_kwargs={}, **kwargs):
        '''
        Perform segmentation on each channel, then cluster the results to find
        bubbles.
        '''

        if verbose:
            output = sys.stdout
        else:
            output = None

        if verbose:
            print("Running bubble finding plane-by-plane.")
        twod_results = \
            ProgressBar.map(_region_return,
                            ((self.cube[i],
                              self.cube.mask.include(view=(i, ))
                              if use_cube_mask else None,
                              i, self.sigma, nsig, overlap_frac,
                              self.keep_threshold_mask, self.distance, scales)
                             for i in xrange(self.cube.shape[0])),
                            multiprocess=multiprocess,
                            nprocesses=nprocesses,
                            file=output,
                            step=self.cube.shape[0],
                            item_len=self.cube.shape[0])

        twod_regions = []
        if self.keep_threshold_mask:
            self._mask = np.zeros(self.cube.shape, dtype=np.bool)

        for out in twod_results:
            if self.keep_threshold_mask:
                chan, regions, mask_slice = out

                self._mask[chan] = mask_slice
            else:
                chan, regions = out

            twod_regions.extend(regions)

        self._bubbles = []
        self._unclustered_regions = []

        if len(twod_regions) == 0:
            warn("No bubbles found in the given cube.")
            return self

        bubble_props = np.vstack([bub.params for bub in twod_regions])

        if verbose:
            print("Clustering 2D regions across channels.")
        # cluster_idx = cluster_and_clean(bubble_props, **kwargs)
        cluster_idx = cluster_brute_force(bubble_props, **kwargs)

        # Add the unclustered ones first
        unclusts = [twod_regions[idx] for idx in np.where(cluster_idx == 0)[0]]
        for reg in unclusts:
            self._unclustered_regions.append([reg])

        good_clusters = []
        for idx in np.unique(cluster_idx[cluster_idx > 0]):
            regions = [twod_regions[idx] for idx in
                       np.where(cluster_idx == idx)[0]]

            if len(regions) < min_channels:
                self._unclustered_regions.append(regions)
                continue

            chans = np.array([reg.channel_center for reg in regions])
            if chans.max() + 1 - chans.min() >= min_channels:
                good_clusters.append(regions)
            else:
                self._unclustered_regions.append(regions)

        if verbose:
            print("Creating bubbles and finding their properties.")
        # We need to pass a linewidth array for the cube, but don't want to
        # have to recompute it multiple times. Make sure it's using the mask
        if cube_linewidth is None:
            # This gives some funky results
            # cube_linewidth = self.cube.with_mask(self.mask).linewidth_fwhm()
            # Just mask with a sigma cut
            sigma_w_unit = self.sigma * self.cube.unit
            cube_linewidth = \
                self.cube.with_mask(self.cube >= 3 *
                                    sigma_w_unit).linewidth_fwhm()
        # Now create the bubble objects and find their respective properties
        self._bubbles = ProgressBar.map(_make_bubble,
                                        ((regions, refit, self.cube, self.mask,
                                          self.distance, self.sigma,
                                          cube_linewidth, self.galaxy_props)
                                         for regions in good_clusters),
                                        multiprocess=False,
                                        nprocesses=nprocesses,
                                        file=output,
                                        step=1,
                                        item_len=len(good_clusters))

        # Now we prune off overlapping bubbles
        self._bubbles, removed_bubbles, new_twoD_clusters = \
            threeD_overlaps(self.bubbles, **overlap_kwargs)

        # Add the 2D regions in the removed bubbles to the unclustered list
        for bub in removed_bubbles:
            self._unclustered_regions.append(bub.twoD_regions)

        # If there is nothing to join, return here
        if len(new_twoD_clusters) == 0:
            return self

        print("Found bubbles to join together.")

        new_bubbles = \
            ProgressBar.map(_make_bubble,
                            ((regions, refit, self.cube, self.mask,
                              self.distance, self.sigma,
                              cube_linewidth, self.galaxy_props)
                             for regions in new_twoD_clusters),
                            multiprocess=False,
                            nprocesses=nprocesses,
                            file=output,
                            step=1,
                            item_len=len(new_twoD_clusters))

        self._bubbles.extend(new_bubbles)

        return self

    @staticmethod
    def reload(cube, bubbles, mask=None, distance=None, galaxy_props=None):
        '''
        Reload from a cube and a list of bubbles.
        '''

        if mask is not None:
            assert cube.shape == mask.shape

        self = BubbleFinder(cube, mask=mask, distance=distance,
                            galaxy_props=galaxy_props)

        self._bubbles = bubbles
        self._unclustered_regions = []

        return self

    @property
    def bubbles(self):
        return self._bubbles

    @property
    def num_bubbles(self):
        return len(self.bubbles)

    @property
    def unclustered_regions(self):
        return self._unclustered_regions

    def to_catalog(self):
        '''
        Returns a PPV_Catalog to explore the population properties.
        '''
        if self.num_bubbles == 0:
            warn("No bubbles were found. Cannot return a catalog.")
            return

        return PPV_Catalog(self.bubbles, self.galaxy_props)

    def visualize_bubbles(self, show=True, edges=False, ax=None,
                          moment0=None, region_col='b', edge_col='g',
                          log_scale=False, plot_threeD_shapes=True,
                          plot_twoD_shapes=False, plot_unclustered=False):
        '''
        Show the location of the bubbles on the moment 0 array.
        '''
        if len(self.bubbles) == 0:
            warn("No bubbles were found. Nothing to show.")
            return

        if moment0 is None:
            # Create the moment array from the cube
            moment0 = self.cube.moment0().value

        import matplotlib.pyplot as p

        if ax is None:
            ax = p.subplot(111)

        if log_scale:
            ax.imshow(np.log10(moment0), cmap='afmhot', origin='lower')
        else:
            ax.imshow(moment0, cmap='afmhot', origin='lower')

        if plot_unclustered:
            for clust in self.unclustered_regions:
                for reg in clust:
                    ax.add_patch(reg.as_patch(color=region_col, fill=False,
                                              linewidth=2))
                    ax.plot(reg.x, reg.y, region_col + 'D')
                    if edges:
                        ax.plot(reg.shell_coords[:, 1], reg.shell_coords[:, 0],
                                edge_col + "o")
        else:
            for bub in self.bubbles:
                if plot_twoD_shapes:
                    for reg in bub.twoD_regions:
                        ax.add_patch(reg.as_patch(color=region_col,
                                                  fill=False,
                                                  linewidth=2))
                        ax.plot(reg.x, reg.y, region_col + "D")
                        if edges:
                            ax.plot(reg.shell_coords[:, 1],
                                    reg.shell_coords[:, 0], edge_col + "o")
                if plot_threeD_shapes:
                    if not plot_twoD_shapes:
                        linewidth = 2
                        linestyle = "-"
                        color = region_col
                    else:
                        linewidth = 3
                        linestyle = "--"
                        color = 'm'
                    ax.add_patch(bub.as_patch(color=color, fill=False,
                                              linewidth=linewidth,
                                              linestyle=linestyle))
                    ax.plot(bub.x, bub.y, region_col + 'D')
                    if edges:
                        ax.plot(bub.shell_coords[:, 2], bub.shell_coords[:, 1],
                                edge_col + "o")

        p.xlim([0, self.cube.shape[2]])
        p.ylim([0, self.cube.shape[1]])

        if show:
            p.show()
        else:
            return ax

    def visualize_channel_maps(self, all_chans=False, subplot=False,
                               edges=False, plot_unclustered=False,
                               interactive_plot=True,
                               show_mask_contours=False, start_chan=None):
        '''
        Plot each channel optionally overlaid with the regions and/or the
        edges.
        '''

        # Loop through all, or just those that have at least one region
        if all_chans:
            chans = np.arange(self.cube.shape[0], dtype=np.int)
        else:
            chans = []
            for bub in self.bubbles:
                chans.extend(list(bub.twoD_region_params()[:, -1]))
            chans = np.unique(np.array(chans, dtype=np.int))

        import matplotlib.pyplot as p
        if interactive_plot:
            p.ion()

        if start_chan is not None:
            if start_chan >= self.cube.shape[0]:
                raise ValueError("start_chan must be below the number of"
                                 " channels: {}".format(self.cube.shape[0]))

            chans = chans[chans >= start_chan]

        for chan in chans:
            if subplot:
                raise NotImplementedError("Need subplots at some point.")
            else:
                ax = p.subplot(111)

            ax.imshow(self.cube[chan].value, cmap='afmhot', origin='lower')

            if show_mask_contours:
                if self.keep_threshold_mask:
                    ax.contour(self.mask[chan], colors='k')
                else:
                    warn("'keep_threshold_mask' must be enabled to plot mask "
                         "contours.")

            for bub in self.bubbles:
                if chan < bub.channel_start and chan > bub.channel_end:
                    continue
                for twod in bub.twoD_regions:
                    if twod.channel_center != chan:
                        continue
                    ax.add_patch(twod.as_patch(color='b', fill=False,
                                               linewidth=2))
                    ax.plot(twod.x, twod.y, 'bD')
                    if edges:
                        ax.plot(twod.shell_coords[:, 1],
                                twod.shell_coords[:, 0], "go")

            # Also show the unclustered regions in each channel
            if plot_unclustered:
                for clust in self.unclustered_regions:
                    for reg in clust:
                        if reg.channel_center != chan:
                            continue
                        ax.add_patch(reg.as_patch(color='r', fill=False,
                                                  linewidth=2))
                        ax.plot(reg.x, reg.y, 'rD')
                        if edges:
                            ax.plot(reg.shell_coords[:, 1],
                                    reg.shell_coords[:, 0], "ro")

            p.xlim([0, self.cube.shape[2]])
            p.ylim([0, self.cube.shape[1]])

            if interactive_plot:
                p.draw()
                raw_input("Channel {}".format(chan))
                p.clf()
            else:
                p.show()

    def save_bubbles(self, folder=None, name=None):
        '''
        Save bubbles as pickled objects.

        Parameters
        ----------
        folder : str, optional
            Path to where objects will be saved.
        name : str, optional
            Prefix for the save names.
        '''

        if len(self.bubbles) == 0:
            warn("There are no bubbles. Returning.")
            return

        import os

        if folder is None:
            folder = ""

        if name is None:
            name = ""

        file_prefix = os.path.join(folder, name)

        for i, bub in enumerate(self.bubbles):
            if len(file_prefix) == 0:
                save_name = "bubble_{}.pkl".format(i)
            else:
                save_name = "{0}_bubble_{1}.pkl".format(file_prefix, i)

            bub.save_bubble(save_name)

def _region_return(imps):
    arr, mask, i, sigma, nsig, overlap_frac, return_mask, distance, scales = \
        imps
    bubs = BubbleFinder2D(arr, channel=i,
                          mask=mask, sigma=sigma, auto_cut=True,
                          scales=scales).\
        multiscale_bubblefind(nsig=nsig,
                              overlap_frac=overlap_frac,
                              distance=distance)
    if return_mask:
        return i, bubs.regions, \
            bubs.insert_in_shape(bubs.mask, bubs._orig_shape, fill_value=True,
                                 dtype=bool)

    return i, bubs.regions

def _make_bubble(imps):
    regions, refit, cube, mask, distance, sigma, lwidth, galaxy_props = imps
    return Bubble3D.from_2D_regions(regions, refit=refit,
                                    cube=cube, mask=mask,
                                    distance=distance,
                                    sigma=sigma, linewidth=lwidth,
                                    galaxy_kwargs=galaxy_props)
<commit_msg>Save channel maps<commit_after>
import numpy as np
import astropy.units as u
from spectral_cube import SpectralCube
# from astropy.utils.console import ProgressBar
import sys
from warnings import warn

from bubble_segment2D import BubbleFinder2D
from bubble_objects import Bubble3D
from bubble_catalog import PPV_Catalog
from clustering import cluster_brute_force, threeD_overlaps
from utils import sig_clip
from galaxy_utils import gal_props_checker
from progressbar import ProgressBar

class BubbleFinder(object):
    """docstring for BubbleFinder"""
    def __init__(self, cube, wcs=None, mask=None, sigma=None, empty_channel=0,
                 keep_threshold_mask=True, distance=None, galaxy_props={}):
        super(BubbleFinder, self).__init__()

        if not isinstance(cube, SpectralCube):
            if wcs is None:
                raise TypeError("When cube is not a SpectralCube, wcs must be"
                                " given.")
            cube = SpectralCube(cube, wcs)
            if mask is not None:
                cube = cube.with_mask(mask)

        self.cube = cube

        self.empty_channel = empty_channel

        if sigma is None:
            self.estimate_sigma()
        else:
            self.sigma = sigma

        self.keep_threshold_mask = keep_threshold_mask
        self._mask = None
        self.distance = distance
        self.galaxy_props = galaxy_props

    @property
    def cube(self):
        return self._cube

    @cube.setter
    def cube(self, input_cube):
        if input_cube.ndim != 3:
            raise TypeError("Given cube must be 3D.")

        self._cube = input_cube

    @property
    def mask(self):
        if not self.keep_threshold_mask:
            warn("Enable keep_threshold_mask to access the entire mask.")

        return self._mask

    def estimate_sigma(self, nsig=10):
        '''
        Use empty channels to estimate sigma. Uses iterative sigma clipping
        to obtain a robust estimate.
        '''

        self.sigma = sig_clip(self.cube[self.empty_channel], nsig=nsig)

    @property
    def sigma(self):
        return self._sigma

    @sigma.setter
    def sigma(self, val):
        if val < 0.0 or ~np.isfinite(val):
            raise ValueError("sigma must be positive and finite.")

        self._sigma = val

    @property
    def galaxy_props(self):
        return self._galaxy_props

    @galaxy_props.setter
    def galaxy_props(self, input_dict):

        # Make sure all of the kwargs are given.
        gal_props_checker(input_dict)

        self._galaxy_props = input_dict

    def get_bubbles(self, verbose=True, overlap_frac=0.9, min_channels=3,
                    use_cube_mask=False, nsig=2., refit=False, scales=None,
                    cube_linewidth=None, multiprocess=True, nprocesses=None,
                    overlap_kwargs={}, **kwargs):
        '''
        Perform segmentation on each channel, then cluster the results to find
        bubbles.
        '''

        if verbose:
            output = sys.stdout
        else:
            output = None

        if verbose:
            print("Running bubble finding plane-by-plane.")
        twod_results = \
            ProgressBar.map(_region_return,
                            ((self.cube[i],
                              self.cube.mask.include(view=(i, ))
                              if use_cube_mask else None,
                              i, self.sigma, nsig, overlap_frac,
                              self.keep_threshold_mask, self.distance, scales)
                             for i in xrange(self.cube.shape[0])),
                            multiprocess=multiprocess,
                            nprocesses=nprocesses,
                            file=output,
                            step=self.cube.shape[0],
                            item_len=self.cube.shape[0])

        twod_regions = []
        if self.keep_threshold_mask:
            self._mask = np.zeros(self.cube.shape, dtype=np.bool)

        for out in twod_results:
            if self.keep_threshold_mask:
                chan, regions, mask_slice = out

                self._mask[chan] = mask_slice
            else:
                chan, regions = out

            twod_regions.extend(regions)

        self._bubbles = []
        self._unclustered_regions = []

        if len(twod_regions) == 0:
            warn("No bubbles found in the given cube.")
            return self

        bubble_props = np.vstack([bub.params for bub in twod_regions])

        if verbose:
            print("Clustering 2D regions across channels.")
        # cluster_idx = cluster_and_clean(bubble_props, **kwargs)
        cluster_idx = cluster_brute_force(bubble_props, **kwargs)

        # Add the unclustered ones first
        unclusts = [twod_regions[idx] for idx in np.where(cluster_idx == 0)[0]]
        for reg in unclusts:
            self._unclustered_regions.append([reg])

        good_clusters = []
        for idx in np.unique(cluster_idx[cluster_idx > 0]):
            regions = [twod_regions[idx] for idx in
                       np.where(cluster_idx == idx)[0]]

            if len(regions) < min_channels:
                self._unclustered_regions.append(regions)
                continue

            chans = np.array([reg.channel_center for reg in regions])
            if chans.max() + 1 - chans.min() >= min_channels:
                good_clusters.append(regions)
            else:
                self._unclustered_regions.append(regions)

        if verbose:
            print("Creating bubbles and finding their properties.")
        # We need to pass a linewidth array for the cube, but don't want to
        # have to recompute it multiple times. Make sure it's using the mask
        if cube_linewidth is None:
            # This gives some funky results
            # cube_linewidth = self.cube.with_mask(self.mask).linewidth_fwhm()
            # Just mask with a sigma cut
            sigma_w_unit = self.sigma * self.cube.unit
            cube_linewidth = \
                self.cube.with_mask(self.cube >= 3 *
                                    sigma_w_unit).linewidth_fwhm()
        # Now create the bubble objects and find their respective properties
        self._bubbles = ProgressBar.map(_make_bubble,
                                        ((regions, refit, self.cube, self.mask,
                                          self.distance, self.sigma,
                                          cube_linewidth, self.galaxy_props)
                                         for regions in good_clusters),
                                        multiprocess=False,
                                        nprocesses=nprocesses,
                                        file=output,
                                        step=1,
                                        item_len=len(good_clusters))

        # Now we prune off overlapping bubbles
        self._bubbles, removed_bubbles, new_twoD_clusters = \
            threeD_overlaps(self.bubbles, **overlap_kwargs)

        # Add the 2D regions in the removed bubbles to the unclustered list
        for bub in removed_bubbles:
            self._unclustered_regions.append(bub.twoD_regions)

        # If there is nothing to join, return here
        if len(new_twoD_clusters) == 0:
            return self

        print("Found bubbles to join together.")

        new_bubbles = \
            ProgressBar.map(_make_bubble,
                            ((regions, refit, self.cube, self.mask,
                              self.distance, self.sigma,
                              cube_linewidth, self.galaxy_props)
                             for regions in new_twoD_clusters),
                            multiprocess=False,
                            nprocesses=nprocesses,
                            file=output,
                            step=1,
                            item_len=len(new_twoD_clusters))

        self._bubbles.extend(new_bubbles)

        return self

    @staticmethod
    def reload(cube, bubbles, mask=None, distance=None, galaxy_props=None):
        '''
        Reload from a cube and a list of bubbles.
        '''

        if mask is not None:
            assert cube.shape == mask.shape

        self = BubbleFinder(cube, mask=mask, distance=distance,
                            galaxy_props=galaxy_props)

        self._bubbles = bubbles
        self._unclustered_regions = []

        return self

    @property
    def bubbles(self):
        return self._bubbles

    @property
    def num_bubbles(self):
        return len(self.bubbles)

    @property
    def unclustered_regions(self):
        return self._unclustered_regions

    def to_catalog(self):
        '''
        Returns a PPV_Catalog to explore the population properties.
        '''
        if self.num_bubbles == 0:
            warn("No bubbles were found. Cannot return a catalog.")
            return

        return PPV_Catalog(self.bubbles, self.galaxy_props)

    def visualize_bubbles(self, show=True, edges=False, ax=None,
                          moment0=None, region_col='b', edge_col='g',
                          log_scale=False, plot_threeD_shapes=True,
                          plot_twoD_shapes=False, plot_unclustered=False):
        '''
        Show the location of the bubbles on the moment 0 array.
        '''
        if len(self.bubbles) == 0:
            warn("No bubbles were found. Nothing to show.")
            return

        if moment0 is None:
            # Create the moment array from the cube
            moment0 = self.cube.moment0().value

        import matplotlib.pyplot as p

        if ax is None:
            ax = p.subplot(111)

        if log_scale:
            ax.imshow(np.log10(moment0), cmap='afmhot', origin='lower')
        else:
            ax.imshow(moment0, cmap='afmhot', origin='lower')

        if plot_unclustered:
            for clust in self.unclustered_regions:
                for reg in clust:
                    ax.add_patch(reg.as_patch(color=region_col, fill=False,
                                              linewidth=2))
                    ax.plot(reg.x, reg.y, region_col + 'D')
                    if edges:
                        ax.plot(reg.shell_coords[:, 1], reg.shell_coords[:, 0],
                                edge_col + "o")
        else:
            for bub in self.bubbles:
                if plot_twoD_shapes:
                    for reg in bub.twoD_regions:
                        ax.add_patch(reg.as_patch(color=region_col,
                                                  fill=False,
                                                  linewidth=2))
                        ax.plot(reg.x, reg.y, region_col + "D")
                        if edges:
                            ax.plot(reg.shell_coords[:, 1],
                                    reg.shell_coords[:, 0], edge_col + "o")
                if plot_threeD_shapes:
                    if not plot_twoD_shapes:
                        linewidth = 2
                        linestyle = "-"
                        color = region_col
                    else:
                        linewidth = 3
                        linestyle = "--"
                        color = 'm'
                    ax.add_patch(bub.as_patch(color=color, fill=False,
                                              linewidth=linewidth,
                                              linestyle=linestyle))
                    ax.plot(bub.x, bub.y, region_col + 'D')
                    if edges:
                        ax.plot(bub.shell_coords[:, 2], bub.shell_coords[:, 1],
                                edge_col + "o")

        p.xlim([0, self.cube.shape[2]])
        p.ylim([0, self.cube.shape[1]])

        if show:
            p.show()
        else:
            return ax

    def visualize_channel_maps(self, all_chans=False, subplot=False,
                               edges=False, plot_unclustered=False,
                               interactive_plot=True,
                               show_mask_contours=False, start_chan=None,
                               save=False, save_path=None, save_name=None):
        '''
        Plot each channel optionally overlaid with the regions and/or the
        edges.
        '''

        # Loop through all, or just those that have at least one region
        if all_chans:
            chans = np.arange(self.cube.shape[0], dtype=np.int)
        else:
            chans = []
            for bub in self.bubbles:
                chans.extend(list(bub.twoD_region_params()[:, -1]))
            chans = np.unique(np.array(chans, dtype=np.int))

        import matplotlib.pyplot as p
        if interactive_plot:
            p.ion()

        if start_chan is not None:
            if start_chan >= self.cube.shape[0]:
                raise ValueError("start_chan must be below the number of"
                                 " channels: {}".format(self.cube.shape[0]))

            chans = chans[chans >= start_chan]

        for chan in chans:
            if subplot:
                raise NotImplementedError("Need subplots at some point.")
            else:
                ax = p.subplot(111)

            ax.imshow(self.cube[chan].value, cmap='afmhot', origin='lower')

            if show_mask_contours:
                if self.keep_threshold_mask:
                    ax.contour(self.mask[chan], colors='k')
                else:
                    warn("'keep_threshold_mask' must be enabled to plot mask "
                         "contours.")

            for bub in self.bubbles:
                if chan < bub.channel_start and chan > bub.channel_end:
                    continue
                for twod in bub.twoD_regions:
                    if twod.channel_center != chan:
                        continue
                    ax.add_patch(twod.as_patch(color='b', fill=False,
                                               linewidth=2))
                    ax.plot(twod.x, twod.y, 'bD')
                    if edges:
                        ax.plot(twod.shell_coords[:, 1],
                                twod.shell_coords[:, 0], "go")

            # Also show the unclustered regions in each channel
            if plot_unclustered:
                for clust in self.unclustered_regions:
                    for reg in clust:
                        if reg.channel_center != chan:
                            continue
                        ax.add_patch(reg.as_patch(color='r', fill=False,
                                                  linewidth=2))
                        ax.plot(reg.x, reg.y, 'rD')
                        if edges:
                            ax.plot(reg.shell_coords[:, 1],
                                    reg.shell_coords[:, 0], "ro")

            p.xlim([0, self.cube.shape[2]])
            p.ylim([0, self.cube.shape[1]])

            if save:
                import os

                if save_path is None:
                    save_path = ""
                if save_name is None:
                    save_name = ""

                p.savefig(os.path.join(save_path,
                                       "{0}_channel_{1}.pdf".format(save_name,
                                                                    chan)))
                p.clf()
            else:
                if interactive_plot:
                    p.draw()
                    raw_input("Channel {}".format(chan))
                    p.clf()
                else:
                    p.show()

    def save_bubbles(self, folder=None, name=None):
        '''
        Save bubbles as pickled objects.

        Parameters
        ----------
        folder : str, optional
            Path to where objects will be saved.
        name : str, optional
            Prefix for the save names.
        '''

        if len(self.bubbles) == 0:
            warn("There are no bubbles. Returning.")
            return

        import os

        if folder is None:
            folder = ""

        if name is None:
            name = ""

        file_prefix = os.path.join(folder, name)

        for i, bub in enumerate(self.bubbles):
            if len(file_prefix) == 0:
                save_name = "bubble_{}.pkl".format(i)
            else:
                save_name = "{0}_bubble_{1}.pkl".format(file_prefix, i)

            bub.save_bubble(save_name)

def _region_return(imps):
    arr, mask, i, sigma, nsig, overlap_frac, return_mask, distance, scales = \
        imps
    bubs = BubbleFinder2D(arr, channel=i,
                          mask=mask, sigma=sigma, auto_cut=True,
                          scales=scales).\
        multiscale_bubblefind(nsig=nsig,
                              overlap_frac=overlap_frac,
                              distance=distance)
    if return_mask:
        return i, bubs.regions, \
            bubs.insert_in_shape(bubs.mask, bubs._orig_shape, fill_value=True,
                                 dtype=bool)

    return i, bubs.regions

def _make_bubble(imps):
    regions, refit, cube, mask, distance, sigma, lwidth, galaxy_props = imps
    return Bubble3D.from_2D_regions(regions, refit=refit,
                                    cube=cube, mask=mask,
                                    distance=distance,
                                    sigma=sigma, linewidth=lwidth,
                                    galaxy_kwargs=galaxy_props)
<|endoftext|>


<commit_before>import mimetypes
import os

from StringIO import StringIO

from django.core.files.base import ContentFile

from .processors import ProcessorPipeline, AutoConvert
from .utils import img_to_fobj, open_image, \
        format_to_extension, extension_to_format, UnknownFormatError, \
        UnknownExtensionError

class SpecFile(ContentFile):
    """
    Wraps a ContentFile in a file-like object with a filename
    and a content_type.
    """
    def __init__(self, filename, content):
        self.file = ContentFile(content)
        self.file.name = filename
        try:
            self.file.content_type = mimetypes.guess_type(filename)[0]
        except IndexError:
            self.file.content_type = None

    def __str__(self):
        return self.file.name

class SpecFileGenerator(object):
    def __init__(self, processors=None, format=None, options={},
            autoconvert=True, storage=None):
        self.processors = processors
        self.format = format
        self.options = options
        self.autoconvert = autoconvert
        self.storage = storage

    def process_content(self, content, filename=None, source_file=None):
        img = open_image(content)
        original_format = img.format

        # Run the processors
        processors = self.processors
        if callable(processors):
            processors = processors(source_file)
        img = ProcessorPipeline(processors or []).process(img)

        options = dict(self.options or {})

        # Determine the format.
        format = self.format
        if not format:
            # Try to guess the format from the extension.
            extension = os.path.splitext(filename)[1].lower()
            if extension:
                try:
                    format = extension_to_format(extension)
                except UnknownExtensionError:
                    pass
        format = format or img.format or original_format or 'JPEG'

        # Run the AutoConvert processor
        if self.autoconvert:
            autoconvert_processor = AutoConvert(format)
            img = autoconvert_processor.process(img)
            options = dict(autoconvert_processor.save_kwargs.items() + \
                    options.items())

        imgfile = img_to_fobj(img, format, **options)
        content = SpecFile(filename, imgfile.read())
        return img, content

    def suggest_extension(self, name):
        original_extension = os.path.splitext(name)[1]
        try:
            suggested_extension = format_to_extension(self.format)
        except UnknownFormatError:
            extension = original_extension
        else:
            if suggested_extension.lower() == original_extension.lower():
                extension = original_extension
            else:
                try:
                    original_format = extension_to_format(original_extension)
                except UnknownExtensionError:
                    extension = suggested_extension
                else:
                    # If the formats match, give precedence to the original extension.
                    if self.format.lower() == original_format.lower():
                        extension = original_extension
                    else:
                        extension = suggested_extension
        return extension

    def generate_file(self, filename, source_file, save=True):
        """
        Generates a new image file by processing the source file and returns
        the content of the result, ready for saving.

        """
        if source_file:  # TODO: Should we error here or something if the source_file doesn't exist?
            # Process the original image file.

            try:
                fp = source_file.storage.open(source_file.name)
            except IOError:
                return
            fp.seek(0)
            fp = StringIO(fp.read())

            img, content = self.process_content(fp, filename, source_file)

            if save:
                storage = self.storage or source_file.storage
                storage.save(filename, content)

            return content
<commit_msg>Add SpecFile.__unicode__<commit_after>import mimetypes
import os

from StringIO import StringIO

from django.core.files.base import ContentFile
from django.utils.encoding import smart_str, smart_unicode

from .processors import ProcessorPipeline, AutoConvert
from .utils import img_to_fobj, open_image, \
        format_to_extension, extension_to_format, UnknownFormatError, \
        UnknownExtensionError

class SpecFile(ContentFile):
    """
    Wraps a ContentFile in a file-like object with a filename
    and a content_type.
    """
    def __init__(self, filename, content):
        self.file = ContentFile(content)
        self.file.name = filename
        try:
            self.file.content_type = mimetypes.guess_type(filename)[0]
        except IndexError:
            self.file.content_type = None

    def __str__(self):
        return smart_str(self.file.name or '')

    def __unicode__(self):
        return smart_unicode(self.file.name or u'')

class SpecFileGenerator(object):
    def __init__(self, processors=None, format=None, options={},
            autoconvert=True, storage=None):
        self.processors = processors
        self.format = format
        self.options = options
        self.autoconvert = autoconvert
        self.storage = storage

    def process_content(self, content, filename=None, source_file=None):
        img = open_image(content)
        original_format = img.format

        # Run the processors
        processors = self.processors
        if callable(processors):
            processors = processors(source_file)
        img = ProcessorPipeline(processors or []).process(img)

        options = dict(self.options or {})

        # Determine the format.
        format = self.format
        if not format:
            # Try to guess the format from the extension.
            extension = os.path.splitext(filename)[1].lower()
            if extension:
                try:
                    format = extension_to_format(extension)
                except UnknownExtensionError:
                    pass
        format = format or img.format or original_format or 'JPEG'

        # Run the AutoConvert processor
        if self.autoconvert:
            autoconvert_processor = AutoConvert(format)
            img = autoconvert_processor.process(img)
            options = dict(autoconvert_processor.save_kwargs.items() + \
                    options.items())

        imgfile = img_to_fobj(img, format, **options)
        content = SpecFile(filename, imgfile.read())
        return img, content

    def suggest_extension(self, name):
        original_extension = os.path.splitext(name)[1]
        try:
            suggested_extension = format_to_extension(self.format)
        except UnknownFormatError:
            extension = original_extension
        else:
            if suggested_extension.lower() == original_extension.lower():
                extension = original_extension
            else:
                try:
                    original_format = extension_to_format(original_extension)
                except UnknownExtensionError:
                    extension = suggested_extension
                else:
                    # If the formats match, give precedence to the original extension.
                    if self.format.lower() == original_format.lower():
                        extension = original_extension
                    else:
                        extension = suggested_extension
        return extension

    def generate_file(self, filename, source_file, save=True):
        """
        Generates a new image file by processing the source file and returns
        the content of the result, ready for saving.

        """
        if source_file:  # TODO: Should we error here or something if the source_file doesn't exist?
            # Process the original image file.

            try:
                fp = source_file.storage.open(source_file.name)
            except IOError:
                return
            fp.seek(0)
            fp = StringIO(fp.read())

            img, content = self.process_content(fp, filename, source_file)

            if save:
                storage = self.storage or source_file.storage
                storage.save(filename, content)

            return content
<|endoftext|>


<commit_before>######### TOKEN CONFIG #########

# some config on how to split a string into tokens
# - set characters to split on?
# - set characters to append at the end of tokens?

######### LABEL CONFIG #########

# these are the labels for tagging tokens
LABELS = [
'PrefixMarital',
'PrefixOther',
'GivenName',
'FirstInitial',
'MiddleName',
'MiddleInitial',
'Surname',
'LastInitial',
'SuffixGenerational',
'SuffixOther',
'Nickname'
]

NULL_LABEL = 'Null'<commit_msg>xml config<commit_after>######### TOKEN CONFIG #########

# some config on how to split a string into tokens
# - set characters to split on?
# - set characters to append at the end of tokens?

######### LABEL CONFIG #########

# these are the labels for tagging tokens
LABELS = [
'PrefixMarital',
'PrefixOther',
'GivenName',
'FirstInitial',
'MiddleName',
'MiddleInitial',
'Surname',
'LastInitial',
'SuffixGenerational',
'SuffixOther',
'Nickname'
]

NULL_LABEL = 'Null'

######## XML CONFIG ############

# this is the tag for each string
PARENT_LABEL = 'Name'

# this is the tag for a group of strings
GROUP_LABEL = 'NameCollection'<|endoftext|>


<commit_before>
import atexit
from queue import Queue
from subprocess import Popen, PIPE
import sys
from threading import Thread
from tkinter import BOTH, END, Frame, Text, Tk, ttk

import psutil

class HarmonbotGUI:

	def __init__(self, master):
		self.master = master
		master.title("Harmonbot")

		self.notebook = ttk.Notebook(master)
		for tab in ("overview", "discord", "discord_listener", "twitch", "telegram"):
			frame = Frame(self.notebook)
			setattr(self, f"{tab}_tab", frame)
			self.notebook.add(frame, text = tab.replace('_', ' ').title())
		self.notebook.pack()

		for overview_frame in ("discord", "discord_listener", "twitch", "telegram"):
			frame = Frame(self.overview_tab)
			setattr(self, f"overview_{overview_frame}_frame", frame)
			text = Text(frame)
			setattr(self, f"overview_{overview_frame}_text", text)
			text.pack()
		self.overview_discord_frame.grid(row = 1, column = 1)
		self.overview_discord_listener_frame.grid(row = 2, column = 1)
		self.overview_twitch_frame.grid(row = 1, column = 2)
		self.overview_telegram_frame.grid(row = 2, column = 2)

		for tab in ("discord", "discord_listener", "twitch", "telegram"):
			notebook_tab = getattr(self, f"{tab}_tab")
			frame = Frame(notebook_tab)
			setattr(self, f"{tab}_frame", frame)
			frame.pack(expand = True, fill = BOTH)
			text = Text(frame)
			setattr(self, f"{tab}_text", text)
			text.pack(expand = True, fill = BOTH)

if __name__ == "__main__":
	root = Tk()
	root.wm_state("zoomed")  # Maximized window
	harmonbot_gui = HarmonbotGUI(root)

	processes = {}
	process_kwargs = {"stdout": PIPE, "stderr": PIPE, "bufsize": 1, "universal_newlines": True}
	processes["discord"] = Popen([sys.executable, "-u", "Harmonbot.py"], cwd = "Discord", **process_kwargs)
	processes["discord_listener"] = Popen(["go", "run", "Harmonbot_Listener.go"], cwd = "Discord", shell = True, **process_kwargs)
	processes["twitch"] = Popen(["pyw", "-3.6", "-u", "Twitch_Harmonbot.py"], cwd = "Twitch", **process_kwargs)
	# TODO: Update to use Python 3.7 executable
	processes["telegram"] = Popen([sys.executable, "-u", "Telegram_Harmonbot.py"], cwd = "Telegram", **process_kwargs)

	def enqueue_output(out, queue):
		for line in iter(out.readline, ""):
			queue.put(line)
		out.close()

	output_queues = {}
	stdout_threads = {}
	stderr_threads = {}
	for name, process in processes.items():
		output_queue = Queue()
		output_queues[name] = output_queue
		# stdout
		stdout_thread = Thread(target = enqueue_output, args = (process.stdout, output_queue))
		stdout_threads[name] = stdout_thread
		stdout_thread.daemon = True
		stdout_thread.start()
		# stderr
		stderr_thread = Thread(target = enqueue_output, args = (process.stderr, output_queue))
		stderr_threads[name] = stderr_thread
		stderr_thread.daemon = True
		stderr_thread.start()
		# TODO: Check stdout and stderr order

	def process_output(name):
		output_queue = output_queues[name]
		while not output_queue.empty():
			line = output_queue.get_nowait()
			for text_name in (f"overview_{name}_text", f"{name}_text"):
				text = getattr(harmonbot_gui, text_name)
				text.insert(END, line)
		root.after(100, process_output, name)  # Every 1/10 sec.

	def check_process_ended(name):
		if processes[name].poll() is None:
			root.after(100, check_process_ended, name)  # Every 1/10 sec.
		else:
			line = f"{name.replace('_', ' ').title()} process ended"
			for text_name in (f"overview_{name}_text", f"{name}_text"):
				text = getattr(harmonbot_gui, text_name)
				text.insert(END, line)

	for process in processes:
		root.after(0, process_output, process)
		root.after(0, check_process_ended, process)

	def cleanup():
		go_process = psutil.Process(processes["discord_listener"].pid)
		for process in go_process.children(recursive = True):
			process.terminate()
		for process in processes.values():
			process.terminate()
		## root.destroy()

	atexit.register(cleanup)

	## root.protocol("WM_DELETE_WINDOW", cleanup)
	root.mainloop()

<commit_msg>[GUI] Dynamically initialize both stdout and stderr threads<commit_after>
import atexit
from queue import Queue
from subprocess import Popen, PIPE
import sys
from threading import Thread
from tkinter import BOTH, END, Frame, Text, Tk, ttk

import psutil

class HarmonbotGUI:

	def __init__(self, master):
		self.master = master
		master.title("Harmonbot")

		self.notebook = ttk.Notebook(master)
		for tab in ("overview", "discord", "discord_listener", "twitch", "telegram"):
			frame = Frame(self.notebook)
			setattr(self, f"{tab}_tab", frame)
			self.notebook.add(frame, text = tab.replace('_', ' ').title())
		self.notebook.pack()

		for overview_frame in ("discord", "discord_listener", "twitch", "telegram"):
			frame = Frame(self.overview_tab)
			setattr(self, f"overview_{overview_frame}_frame", frame)
			text = Text(frame)
			setattr(self, f"overview_{overview_frame}_text", text)
			text.pack()
		self.overview_discord_frame.grid(row = 1, column = 1)
		self.overview_discord_listener_frame.grid(row = 2, column = 1)
		self.overview_twitch_frame.grid(row = 1, column = 2)
		self.overview_telegram_frame.grid(row = 2, column = 2)

		for tab in ("discord", "discord_listener", "twitch", "telegram"):
			notebook_tab = getattr(self, f"{tab}_tab")
			frame = Frame(notebook_tab)
			setattr(self, f"{tab}_frame", frame)
			frame.pack(expand = True, fill = BOTH)
			text = Text(frame)
			setattr(self, f"{tab}_text", text)
			text.pack(expand = True, fill = BOTH)

if __name__ == "__main__":
	root = Tk()
	root.wm_state("zoomed")  # Maximized window
	harmonbot_gui = HarmonbotGUI(root)

	processes = {}
	process_kwargs = {"stdout": PIPE, "stderr": PIPE, "bufsize": 1, "universal_newlines": True}
	processes["discord"] = Popen([sys.executable, "-u", "Harmonbot.py"], cwd = "Discord", **process_kwargs)
	processes["discord_listener"] = Popen(["go", "run", "Harmonbot_Listener.go"], cwd = "Discord", shell = True, **process_kwargs)
	processes["twitch"] = Popen(["pyw", "-3.6", "-u", "Twitch_Harmonbot.py"], cwd = "Twitch", **process_kwargs)
	# TODO: Update to use Python 3.7 executable
	processes["telegram"] = Popen([sys.executable, "-u", "Telegram_Harmonbot.py"], cwd = "Telegram", **process_kwargs)

	def enqueue_output(out, queue):
		for line in iter(out.readline, ""):
			queue.put(line)
		out.close()

	output_queues = {}
	output_threads = {"stdout": {}, "stderr": {}}
	for name, process in processes.items():
		output_queue = Queue()
		output_queues[name] = output_queue
		for output_type in ("stdout", "stderr"):
			process_output = getattr(process, output_type)
			output_thread = Thread(target = enqueue_output, args = (process_output, output_queue))
			output_threads[output_type][name] = output_thread
			output_thread.daemon = True
			output_thread.start()

	# TODO: Check stdout and stderr order

	def process_output(name):
		output_queue = output_queues[name]
		while not output_queue.empty():
			line = output_queue.get_nowait()
			for text_name in (f"overview_{name}_text", f"{name}_text"):
				text = getattr(harmonbot_gui, text_name)
				text.insert(END, line)
		root.after(100, process_output, name)  # Every 1/10 sec.

	def check_process_ended(name):
		if processes[name].poll() is None:
			root.after(100, check_process_ended, name)  # Every 1/10 sec.
		else:
			line = f"{name.replace('_', ' ').title()} process ended"
			for text_name in (f"overview_{name}_text", f"{name}_text"):
				text = getattr(harmonbot_gui, text_name)
				text.insert(END, line)

	for process in processes:
		root.after(0, process_output, process)
		root.after(0, check_process_ended, process)

	def cleanup():
		go_process = psutil.Process(processes["discord_listener"].pid)
		for process in go_process.children(recursive = True):
			process.terminate()
		for process in processes.values():
			process.terminate()
		## root.destroy()

	atexit.register(cleanup)

	## root.protocol("WM_DELETE_WINDOW", cleanup)
	root.mainloop()

<|endoftext|>