<commit_before>#!/usr/bin/env python3
# Copyright (C) 2020 The Android Open Source Project
#
# 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 os

# This function lives in its own module to allow us to use a different
# mechanism of loading depending on the build system.
def read_descriptor():
  ws = os.path.dirname(__file__)
  with open(os.path.join(ws, 'trace_processor.descriptor'), 'rb') as x:
    return x.read()
<commit_msg>Add try block to read_descriptor for it to read the descriptor file correctly on Colab<commit_after>#!/usr/bin/env python3
# Copyright (C) 2020 The Android Open Source Project
#
# 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 os

# This function lives in its own module to allow us to use a different
# mechanism of loading depending on the build system.
def read_descriptor():
  try:
    from .descriptor_vendor import read_tp_descriptor
    return read_tp_descriptor()
  except ModuleNotFoundError:
    ws = os.path.dirname(__file__)
    with open(os.path.join(ws, 'trace_processor.descriptor'), 'rb') as x:
      return x.read()
<|endoftext|>


<commit_before>from django.contrib import admin
from django.utils.timesince import timesince

from . import models

@admin.register(models.DomainCheck)
class DomainCheckAdmin(admin.ModelAdmin):

    list_display = (
        'domain', 'path', 'protocol', 'method', 'is_active',
        'status', 'last_checked', )
    list_filter = ('protocol', 'method', 'is_active', )
    search_fields = ('domain', )

    def get_queryset(self, request):
        return super().get_queryset(request).status()

    def status(self, obj):
        return obj.status.title()
    status.admin_order_field = 'success_rate'

    def last_checked(self, obj):
        if obj.last_check:
            return '{} ago'.format(timesince(obj.last_check))
        else:
            return 'Never'
    last_checked.admin_order_field = 'last_check'

@admin.register(models.CheckResult)
class CheckResultAdmin(admin.ModelAdmin):

    date_hierarchy = 'checked_on'
    list_display = ('domain_check', 'status_code', )
    list_filter = ('checked_on', )
<commit_msg>Adding the StatusListFilter for the DomainChecks.<commit_after>from django.contrib import admin
from django.utils.timesince import timesince

from . import models

class StatusListFilter(admin.SimpleListFilter):
    title = 'status'
    parameter_name = 'status'

    def lookups(self, request, model_admin):
        return (
            ('good', 'Good'),
            ('fair', 'Fair'),
            ('poor', 'Poor'),
            ('unknown', 'Unknown'),
        )

    def queryset(self, request, queryset):
        if self.value():
            return queryset.filter(status=self.value())

@admin.register(models.DomainCheck)
class DomainCheckAdmin(admin.ModelAdmin):

    list_display = (
        'domain', 'path', 'protocol', 'method', 'is_active',
        'status', 'last_checked', )
    list_filter = ('protocol', 'method', StatusListFilter, 'is_active', )
    search_fields = ('domain', )

    def get_queryset(self, request):
        return super().get_queryset(request).status()

    def status(self, obj):
        return obj.status.title()
    status.admin_order_field = 'success_rate'

    def last_checked(self, obj):
        if obj.last_check:
            return '{} ago'.format(timesince(obj.last_check))
        else:
            return 'Never'
    last_checked.admin_order_field = 'last_check'

@admin.register(models.CheckResult)
class CheckResultAdmin(admin.ModelAdmin):

    date_hierarchy = 'checked_on'
    list_display = ('domain_check', 'status_code', )
    list_filter = ('checked_on', )
<|endoftext|>


<commit_before>from hashlib import sha1
import couchdb
from core import Store

# Add time serialization to couchdb's json repertoire.
import json
import time
import calendar
class TimeEncoder(json.JSONEncoder):
    def default(self, obj):
        if type(obj) is time.struct_time:
            return calendar.timegm(obj)
        else:
            return json.JSONEncoder.default(self, obj)
couchdb.json.use(
        decode=json.JSONDecoder().decode,
        encode=TimeEncoder().encode)

class CouchStore(Store):
    def load(self):
        self.server = couchdb.client.Server()

        db_name = self.config.get('db_name', 'activity')
        if not db_name in self.server:
            self.server.create(db_name)
        self.db = self.server[db_name]
        self.init_views()

        return self

    def init_views(self):
        views = {
            "_id": "_design/activity",
            "language": "javascript",
            "views": {
                "sources": {
                    "map": "function(doc) { if (doc.type == 'source') { emit(doc._id, doc) } }"
                },
                "events": {
                    "map": "function(doc) { if (doc.type == 'event') { emit(doc.timestamp, doc) } }"
                },
                "events-by-source": {
                    "map": "function(doc) { if (doc.type == 'event') { emit(doc.source, doc) } }"
                },
                "events-by-source-count": {
                    "map": "function(doc) { if (doc.type == 'event') { emit(doc.source, doc) } }",
                    "reduce": "_count"
                },
                "events-by-slug": {
                    "map": "function(doc) { if (doc.type == 'event' && doc.slug) { emit(doc.slug, doc) } }"
                }
            }
        }
        doc = self.db.get(views['_id'], {})
        doc.update(views)
        self.db[views['_id']] = doc

    def get_sources(self):
        return dict((row.key, row.value) for row in self.db.view('activity/sources'))

    def store_source(self, source):
        source_data = source.config.copy()
        source_data['type'] = 'source'
        self.db[source.source_id] = source_data

    def store_events(self, events):
        ids = {}
        for event in events:
            event.setdefault('_id', sha1(event['summary'].encode('utf-8')+str(event['timestamp'])).hexdigest())
            ids[event['_id']] = event
            event['type'] = 'event'

        tries = 3
        while ids and tries:
            tries -= 1
            result = self.db.update(ids.values())
            for success, _id, info in result:
                if success:
                    del ids[_id]
                else:
                    ids[_id]['_rev'] = self.db[_id]['_rev']

        if ids:
            raise couchdb.ResourceConflict

    def store_update(self, source, events):
        for event in events:
            event['kind'] = source.kind
            event['source'] = source.source_id
        self.store_events(events)
        self.db[source.source_id] = source.config

    def events(self, count=100):
        return (event.value for event in self.db.view('activity/events', limit=count, descending=True))

    def events_by_slug(self, slug):
        return (event.value for event in self.db.view('activity/events-by-slug')[slug])

    def events_by_source_count(self):
        return dict((count.key, count.value) for count in self.db.view('activity/events-by-source-count', group_level=1))

    def empty(self):
        for event in self.db.view('activity/events'):
            self.db.delete(event.value)
        for row in self.db.view('activity/sources'):
            source = row.value
            source['since'] = {}
            self.db[row.id] = source
<commit_msg>Allow startkey to be specified, enabling pagination.<commit_after>from hashlib import sha1
import couchdb
from core import Store

# Add time serialization to couchdb's json repertoire.
import json
import time
import calendar
class TimeEncoder(json.JSONEncoder):
    def default(self, obj):
        if type(obj) is time.struct_time:
            return calendar.timegm(obj)
        else:
            return json.JSONEncoder.default(self, obj)
couchdb.json.use(
        decode=json.JSONDecoder().decode,
        encode=TimeEncoder().encode)

class CouchStore(Store):
    def load(self):
        self.server = couchdb.client.Server()

        db_name = self.config.get('db_name', 'activity')
        if not db_name in self.server:
            self.server.create(db_name)
        self.db = self.server[db_name]
        self.init_views()

        return self

    def init_views(self):
        views = {
            "_id": "_design/activity",
            "language": "javascript",
            "views": {
                "sources": {
                    "map": "function(doc) { if (doc.type == 'source') { emit(doc._id, doc) } }"
                },
                "events": {
                    "map": "function(doc) { if (doc.type == 'event') { emit(doc.timestamp, doc) } }"
                },
                "events-by-source": {
                    "map": "function(doc) { if (doc.type == 'event') { emit(doc.source, doc) } }"
                },
                "events-by-source-count": {
                    "map": "function(doc) { if (doc.type == 'event') { emit(doc.source, doc) } }",
                    "reduce": "_count"
                },
                "events-by-slug": {
                    "map": "function(doc) { if (doc.type == 'event' && doc.slug) { emit(doc.slug, doc) } }"
                }
            }
        }
        doc = self.db.get(views['_id'], {})
        doc.update(views)
        self.db[views['_id']] = doc

    def get_sources(self):
        return dict((row.key, row.value) for row in self.db.view('activity/sources'))

    def store_source(self, source):
        source_data = source.config.copy()
        source_data['type'] = 'source'
        self.db[source.source_id] = source_data

    def store_events(self, events):
        ids = {}
        for event in events:
            event.setdefault('_id', sha1(event['summary'].encode('utf-8')+str(event['timestamp'])).hexdigest())
            ids[event['_id']] = event
            event['type'] = 'event'

        tries = 3
        while ids and tries:
            tries -= 1
            result = self.db.update(ids.values())
            for success, _id, info in result:
                if success:
                    del ids[_id]
                else:
                    ids[_id]['_rev'] = self.db[_id]['_rev']

        if ids:
            raise couchdb.ResourceConflict

    def store_update(self, source, events):
        for event in events:
            event['kind'] = source.kind
            event['source'] = source.source_id
        self.store_events(events)
        self.db[source.source_id] = source.config

    def events(self, count=100, before=None):
        options = { 'limit': count, 'descending': True }
        if before is not None:
            options['startkey'] = before
        return (event.value for event in self.db.view('activity/events', **options))

    def events_by_slug(self, slug):
        return (event.value for event in self.db.view('activity/events-by-slug')[slug])

    def events_by_source_count(self):
        return dict((count.key, count.value) for count in self.db.view('activity/events-by-source-count', group_level=1))

    def empty(self):
        for event in self.db.view('activity/events'):
            self.db.delete(event.value)
        for row in self.db.view('activity/sources'):
            source = row.value
            source['since'] = {}
            self.db[row.id] = source
<|endoftext|>


<commit_before>import json
import sys
import time
from datetime import datetime
from os import path

import xbmcaddon
import xbmcgui
import xbmcplugin

class Menu(object):
    _plugin_url = sys.argv[0]
    _handle = int(sys.argv[1])
    _addon_path = None

    def __init__(self, sort_methods=[]):
        addon = xbmcaddon.Addon("plugin.video.nfl-teams")
        self._addon_path = addon.getAddonInfo("path")

        for method in sort_methods:
            self._add_sort_method(method)

    def _add_sort_method(self, sort_method="none"):
        if sort_method == "none":
            xbmcplugin.addSortMethod(self._handle, xbmcplugin.SORT_METHOD_NONE)
        elif sort_method == "alpha":
            xbmcplugin.addSortMethod(self._handle, xbmcplugin.SORT_METHOD_LABEL)
            xbmcplugin.addSortMethod(self._handle, xbmcplugin.SORT_METHOD_LABEL_IGNORE_THE)
        elif sort_method == "date":
            xbmcplugin.addSortMethod(self._handle, xbmcplugin.SORT_METHOD_DATE)

    def add_item(self, item):
        url = "{0}?{1}".format(self._plugin_url, json.dumps(item.get("url_params")))

        listitem = xbmcgui.ListItem()
        listitem.setLabel(item.get("name"))

        thumbnail = item.get("thumbnail")
        if thumbnail.startswith("resources"):
            thumbnail = path.join(self._addon_path, thumbnail)
        listitem.setThumbnailImage(thumbnail)

        fanart = path.join(self._addon_path, "resources", "images", "fanart", "{0}.jpg".format(item.get("url_params")["team"]))
        listitem.setProperty("fanart_image", fanart)

        info = {"title": item.get("name")}

        if item.get("raw_metadata"):
            info["plot"] = item.get("raw_metadata").get("description", "No description was given for the video.")
            date = self.parse_video_date(item.get("raw_metadata").get("date"))
            if date:
                info["date"] = date.strftime("%d.%m.%Y")
                info["plot"] = "Added on {0}.\n{1}".format(date.strftime("%c"), info["plot"])

        listitem.setInfo("video", info)

        if item.get("folder"):
            xbmcplugin.addDirectoryItem(self._handle, url, listitem, isFolder=item.get("folder"))
        else:
            xbmcplugin.addDirectoryItem(self._handle, url, listitem)

    @classmethod
    def parse_video_date(cls, date_string):
        if date_string:
            try:
                return datetime.strptime(date_string, "%m/%d/%Y %H:%M:%S")
            except TypeError:
                # Workaround for bug in XBMC: http://forum.xbmc.org/showthread.php?tid=112916
                return datetime.fromtimestamp(time.mktime(time.strptime(date_string, "%m/%d/%Y %H:%M:%S")))
        else:
            return None

    def _end_directory(self):
        xbmcplugin.endOfDirectory(self._handle)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._end_directory()
        return False
<commit_msg>Minor cleanup in Menu class<commit_after>import json
import sys
import time
from datetime import datetime
from os import path

import xbmcaddon
import xbmcgui
import xbmcplugin

class Menu(object):
    _plugin_url = sys.argv[0]
    _handle = int(sys.argv[1])
    _addon_path = None

    def __init__(self, sort_methods):
        addon = xbmcaddon.Addon("plugin.video.nfl-teams")
        self._addon_path = addon.getAddonInfo("path")

        for method in sort_methods:
            self._add_sort_method(method)

    def _add_sort_method(self, sort_method="none"):
        if sort_method == "none":
            xbmcplugin.addSortMethod(self._handle, xbmcplugin.SORT_METHOD_NONE)
        elif sort_method == "alpha":
            xbmcplugin.addSortMethod(self._handle, xbmcplugin.SORT_METHOD_LABEL)
            xbmcplugin.addSortMethod(self._handle, xbmcplugin.SORT_METHOD_LABEL_IGNORE_THE)
        elif sort_method == "date":
            xbmcplugin.addSortMethod(self._handle, xbmcplugin.SORT_METHOD_DATE)

    def add_item(self, item):
        listitem = xbmcgui.ListItem()
        listitem.setLabel(item.get("name"))

        thumbnail = item.get("thumbnail")
        if thumbnail.startswith("resources"):
            thumbnail = path.join(self._addon_path, thumbnail)
        listitem.setThumbnailImage(thumbnail)

        fanart = path.join(self._addon_path, "resources", "images", "fanart", "{0}.jpg".format(item.get("url_params")["team"]))
        listitem.setProperty("fanart_image", fanart)

        info = {"title": item.get("name")}

        if item.get("raw_metadata"):
            info["plot"] = item.get("raw_metadata").get("description", "No description was given for the video.")
            date = self.parse_video_date(item.get("raw_metadata").get("date"))
            if date:
                info["date"] = date.strftime("%d.%m.%Y")
                info["plot"] = "Added on {0}.\n{1}".format(date.strftime("%c"), info["plot"])

        listitem.setInfo("video", info)

        url = "{0}?{1}".format(self._plugin_url, json.dumps(item.get("url_params")))
        if item.get("folder"):
            xbmcplugin.addDirectoryItem(self._handle, url, listitem, isFolder=item.get("folder"))
        else:
            xbmcplugin.addDirectoryItem(self._handle, url, listitem)

    @classmethod
    def parse_video_date(cls, date_string):
        if date_string:
            try:
                return datetime.strptime(date_string, "%m/%d/%Y %H:%M:%S")
            except TypeError:
                # Workaround for bug in XBMC: http://forum.xbmc.org/showthread.php?tid=112916
                return datetime.fromtimestamp(time.mktime(time.strptime(date_string, "%m/%d/%Y %H:%M:%S")))
        else:
            return None

    def _end_directory(self):
        xbmcplugin.endOfDirectory(self._handle)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._end_directory()
        return False
<|endoftext|>


<commit_before>import functools
import itertools
from collections import defaultdict
import time
import logging as log
from Queue import Queue, Empty
import weakref
import threading

from kafka import base
from kafka.common import OffsetType

from .protocol import (PartitionFetchRequest, PartitionOffsetCommitRequest,
                       PartitionOffsetFetchRequest)

class SimpleConsumer(base.BaseSimpleConsumer):
    def __init__(self,
                 topic,
                 cluster,
                 consumer_group=None,
                 partitions=None,
                 fetch_message_max_bytes=1024 * 1024,
                 num_consumer_fetchers=1,
                 auto_commit_enable=False,
                 auto_commit_interval_ms=60 * 1000,
                 queued_max_messages=2000,
                 fetch_min_bytes=1,
                 fetch_wait_max_ms=100,
                 refresh_leader_backoff_ms=200,
                 offsets_channel_backoff_ms=1000,
                 offsets_commit_max_retries=5,
                 auto_offset_reset=OffsetType.LATEST,
                 consumer_timeout_ms=-1):
        """Create a SimpleConsumer.

        Settings and default values are taken from the Scala
        consumer implementation.  Consumer group is included
        because it's necessary for offset management, but doesn't imply
        that this is a balancing consumer. Use a BalancedConsumer for
        that.

        :param topic: the topic this consumer should consume
        :type topic: pykafka.topic.Topic
        :param cluster: the cluster this consumer should connect to
        :type cluster: pykafka.cluster.Cluster
        :param consumer_group: the name of the consumer group to join
        :type consumer_group: str
        :param partitions: existing partitions to which to connect
        :type partitions: list of pykafka.partition.Partition
        :param fetch_message_max_bytes: the number of bytes of messages to
            attempt to fetch
        :type fetch_message_max_bytes: int
        :param num_consumer_fetchers: the number of threads used to fetch data
        :type num_consumer_fetchers: int
        :param auto_commit_enable: if true, periodically commit to kafka the
            offset of messages already fetched by this consumer
        :type auto_commit_enable: bool
        :param auto_commit_interval_ms: the frequency in ms that the consumer
            offsets are committed to kafka
        :type auto_commit_interval_ms: int
        :param queued_max_messages: max number of messages buffered for
            consumption
        :type queued_max_messages: int
        :param fetch_min_bytes: the minimum amount of data the server should
            return for a fetch request. If insufficient data is available the
            request will block
        :type fetch_min_bytes: int
        :param fetch_wait_max_ms: the maximum amount of time the server will
            block before answering the fetch request if there isn't sufficient
            data to immediately satisfy fetch_min_bytes
        :type fetch_wait_max_ms: int
        :param refresh_leader_backoff_ms: backoff time to refresh the leader of
            a partition after it loses the current leader
        :type refresh_leader_backoff_ms: int
        :param offsets_channel_backoff_ms: backoff time to retry offset
            commits/fetches
        :type offsets_channel_backoff_ms: int
        :param offsets_commit_max_retries: Retry the offset commit up to this
            many times on failure.
        :type offsets_commit_max_retries: int
        :param auto_offset_reset: what to do if an offset is out of range
        :type auto_offset_reset: int
        :param consumer_timeout_ms: throw a timeout exception to the consumer
            if no message is available for consumption after the specified interval
        :type consumer_timeout_ms: int
        """
        if not isinstance(cluster, weakref.ProxyType):
            self._cluster = weakref.proxy(cluster)
        else:
            self._cluster = cluster
        self._consumer_group = consumer_group
        self._topic = topic
        self._fetch_message_max_bytes = fetch_message_max_bytes
        self._fetch_min_bytes = fetch_min_bytes
        self._queued_max_messages = queued_max_messages
        self._num_consumer_fetchers = num_consumer_fetchers
        self._fetch_wait_max_ms = fetch_wait_max_ms
        self._consumer_timeout_ms = consumer_timeout_ms

        self._last_message_time = time.time()

        self._auto_commit_enable = auto_commit_enable
        self._auto_commit_interval_ms = auto_commit_interval_ms
        self._last_auto_commit = time.time()

        self._offset_manager = self._cluster.get_offset_manager(self._consumer_group)

        owned_partition_partial = functools.partial(
            OwnedPartition, consumer_group=self._consumer_group)
        if partitions:
            self._partitions = {owned_partition_partial(p): p
                                for p in partitions}
        else:
            self._partitions = {owned_partition_partial(p): topic.partitions[k]
                                for k, p in topic.partitions.iteritems()}
        self._partitions_by_id = {p.partition.id: p
                                  for p in self._partitions.iterkeys()}
        # Organize partitions by leader for efficient queries
        self._partitions_by_leader = defaultdict(list)
        for p in self._partitions.iterkeys():
            self._partitions_by_leader[p.partition.leader].append(p)
        self.partition_cycle = itertools.cycle(self._partitions.keys())

        self._running = True

        if self._auto_commit_enable:
            self._autocommit_worker_thread = self._setup_autocommit_worker()
        # we need to get the most up-to-date offsets before starting consumption
        self.fetch_offsets()
        self._fetch_workers = self._setup_fetch_workers()

    @property
    def topic(self):
        return self._topic

    @property
    def partitions(self):
        return self._partitions

    def __del__(self):
        self.stop()

    def stop(self):
        self._running = False

    def _setup_autocommit_worker(self):
        def autocommitter():
            while True:
                if not self._running:
                    break
                if self._auto_commit_enable:
                    self._auto_commit()
        log.debug("Starting autocommitter thread")
        return self._cluster.handler.spawn(autocommitter)

    def _setup_fetch_workers(self):
        def fetcher():
            while True:
                if not self._running:
                    break
                self.fetch()
        return [self._cluster.handler.spawn(fetcher)
                for i in xrange(self._num_consumer_fetchers)]

    def __iter__(self):
        while True:
            message = self.consume()
            if not message and self._consumer_timed_out():
                raise StopIteration
            yield message

    def _consumer_timed_out(self):
        if self._consumer_timeout_ms == -1:
            return False
        disp = (time.time() - self._last_message_time) * 1000.0
        return disp > self._consumer_timeout_ms

    def consume(self):
        """Get one message from the consumer.

        :param timeout: Seconds to wait before returning None
        """
        owned_partition = self.partition_cycle.next()
        message = owned_partition.consume()

        if message:
            self._last_message_time = time.time()
            return message

    def _auto_commit(self):
        if not self._auto_commit_enable or self._auto_commit_interval_ms == 0:
            return

        if (time.time() - self._last_auto_commit) * 1000.0 >= self._auto_commit_interval_ms:
            log.info("Autocommitting consumer offset for consumer group %s and topic %s",
                     self._consumer_group, self._topic.name)
            self.commit_offsets()
            self._last_auto_commit = time.time()

    def commit_offsets(self):
        """Commit offsets for this consumer's topic

        Uses the offset commit/fetch API
        """
        if not self._consumer_group:
            raise Exception("consumer group must be specified to commit offsets")

        self._offset_manager = self._cluster.get_offset_manager(self._consumer_group)

        reqs = [p.build_offset_commit_request() for p in self._partitions.keys()]
        log.info("Committing offsets for %d partitions", len(reqs))
        self._offset_manager.commit_consumer_group_offsets(
            self._consumer_group, 1, 'pykafka', reqs)

    def fetch_offsets(self):
        """Fetch offsets for this consumer's topic

        Uses the offset commit/fetch API
        Should be called when consumer starts and after any errors
        """
        if not self._consumer_group:
            raise Exception("consumer group must be specified to fetch offsets")

        self._offset_manager = self._cluster.get_offset_manager(self._consumer_group)

        log.info("Fetching offsets")

        reqs = [p.build_offset_fetch_request() for p in self._partitions.keys()]
        res = self._offset_manager.fetch_consumer_group_offsets(self._consumer_group, reqs)
        for partition_id, pres in res.topics[self._topic.name].iteritems():
            partition = self._partitions_by_id[partition_id]
            partition.set_offset_counters(pres)

    def fetch(self):
        """Fetch new messages for all partitions

        Create a FetchRequest for each broker and send it. Enqueue each of the
        returned messages in the approprate OwnedPartition.
        """
        for broker, owned_partitions in self._partitions_by_leader.iteritems():
            reqs = []
            for owned_partition in owned_partitions:
                has_room = owned_partition.message_count < self._queued_max_messages
                if owned_partition.empty and has_room:
                    # attempt to acquire lock, just pass if we can't
                    if owned_partition.lock.acquire(False):
                        reqs.append(owned_partition.build_fetch_request(
                            self._fetch_message_max_bytes))
            if reqs:
                response = broker.fetch_messages(
                    reqs, timeout=self._fetch_wait_max_ms,
                    min_bytes=self._fetch_min_bytes
                )
                for partition_id, pres in response.topics[self._topic.name].iteritems():
                    owned_partition = self._partitions_by_id[partition_id]
                    owned_partition.enqueue_messages(pres.messages)
                    owned_partition.lock.release()

class OwnedPartition(object):
    """A partition that is owned by a SimpleConsumer.

    Used to keep track of offsets and the internal message queue.
    """

    def __init__(self,
                 partition,
                 consumer_group=None):
        self.partition = partition
        self._messages = Queue()
        self.last_offset_consumed = 0
        self.next_offset = 0
        self.lock = threading.Lock()

    @property
    def message_count(self):
        return self._messages.qsize()

    @property
    def empty(self):
        return self._messages.empty()

    def set_offset_counters(self, res):
        """Set the internal offset counters from an OffsetFetchResponse

        :param res: an OffsetFetchPartitionResponse containing the committed
            offset for this partition
        :type res: <protocol.PartitionOffsetFetchResponse>
        """
        self.last_offset_consumed = res.offset
        self.next_offset = res.offset + 1
        log.info("Last offset consumed: %d", self.last_offset_consumed)

    def build_fetch_request(self, max_bytes):
        return PartitionFetchRequest(
            self.partition.topic.name, self.partition.id,
            self.next_offset, max_bytes)

    def build_offset_commit_request(self):
        """Create a PartitionOffsetCommitRequest for this partition
        """
        return PartitionOffsetCommitRequest(
            self.partition.topic.name,
            self.partition.id,
            self.last_offset_consumed,
            int(time.time()),
            ''  # TODO - what to do with metadata?
        )

    def build_offset_fetch_request(self):
        """Create a PartitionOffsetFetchRequest for this partition
        """
        return PartitionOffsetFetchRequest(
            self.partition.topic.name,
            self.partition.id
        )

    def consume(self):
        """Get a single message from this partition
        """
        try:
            message = self._messages.get_nowait()
            self.last_offset_consumed = message.offset
            return message
        except Empty:
            return None

    def enqueue_messages(self, messages):
        for message in messages:
            if message.offset < self.last_offset_consumed:
                continue
            self._messages.put(message)
            self.next_offset = message.offset + 1
<commit_msg>set metadata<commit_after>import functools
import itertools
from collections import defaultdict
import time
import logging as log
from Queue import Queue, Empty
import weakref
import threading

from kafka import base
from kafka.common import OffsetType

from .protocol import (PartitionFetchRequest, PartitionOffsetCommitRequest,
                       PartitionOffsetFetchRequest)

class SimpleConsumer(base.BaseSimpleConsumer):
    def __init__(self,
                 topic,
                 cluster,
                 consumer_group=None,
                 partitions=None,
                 fetch_message_max_bytes=1024 * 1024,
                 num_consumer_fetchers=1,
                 auto_commit_enable=False,
                 auto_commit_interval_ms=60 * 1000,
                 queued_max_messages=2000,
                 fetch_min_bytes=1,
                 fetch_wait_max_ms=100,
                 refresh_leader_backoff_ms=200,
                 offsets_channel_backoff_ms=1000,
                 offsets_commit_max_retries=5,
                 auto_offset_reset=OffsetType.LATEST,
                 consumer_timeout_ms=-1):
        """Create a SimpleConsumer.

        Settings and default values are taken from the Scala
        consumer implementation.  Consumer group is included
        because it's necessary for offset management, but doesn't imply
        that this is a balancing consumer. Use a BalancedConsumer for
        that.

        :param topic: the topic this consumer should consume
        :type topic: pykafka.topic.Topic
        :param cluster: the cluster this consumer should connect to
        :type cluster: pykafka.cluster.Cluster
        :param consumer_group: the name of the consumer group to join
        :type consumer_group: str
        :param partitions: existing partitions to which to connect
        :type partitions: list of pykafka.partition.Partition
        :param fetch_message_max_bytes: the number of bytes of messages to
            attempt to fetch
        :type fetch_message_max_bytes: int
        :param num_consumer_fetchers: the number of threads used to fetch data
        :type num_consumer_fetchers: int
        :param auto_commit_enable: if true, periodically commit to kafka the
            offset of messages already fetched by this consumer
        :type auto_commit_enable: bool
        :param auto_commit_interval_ms: the frequency in ms that the consumer
            offsets are committed to kafka
        :type auto_commit_interval_ms: int
        :param queued_max_messages: max number of messages buffered for
            consumption
        :type queued_max_messages: int
        :param fetch_min_bytes: the minimum amount of data the server should
            return for a fetch request. If insufficient data is available the
            request will block
        :type fetch_min_bytes: int
        :param fetch_wait_max_ms: the maximum amount of time the server will
            block before answering the fetch request if there isn't sufficient
            data to immediately satisfy fetch_min_bytes
        :type fetch_wait_max_ms: int
        :param refresh_leader_backoff_ms: backoff time to refresh the leader of
            a partition after it loses the current leader
        :type refresh_leader_backoff_ms: int
        :param offsets_channel_backoff_ms: backoff time to retry offset
            commits/fetches
        :type offsets_channel_backoff_ms: int
        :param offsets_commit_max_retries: Retry the offset commit up to this
            many times on failure.
        :type offsets_commit_max_retries: int
        :param auto_offset_reset: what to do if an offset is out of range
        :type auto_offset_reset: int
        :param consumer_timeout_ms: throw a timeout exception to the consumer
            if no message is available for consumption after the specified interval
        :type consumer_timeout_ms: int
        """
        if not isinstance(cluster, weakref.ProxyType):
            self._cluster = weakref.proxy(cluster)
        else:
            self._cluster = cluster
        self._consumer_group = consumer_group
        self._topic = topic
        self._fetch_message_max_bytes = fetch_message_max_bytes
        self._fetch_min_bytes = fetch_min_bytes
        self._queued_max_messages = queued_max_messages
        self._num_consumer_fetchers = num_consumer_fetchers
        self._fetch_wait_max_ms = fetch_wait_max_ms
        self._consumer_timeout_ms = consumer_timeout_ms

        self._last_message_time = time.time()

        self._auto_commit_enable = auto_commit_enable
        self._auto_commit_interval_ms = auto_commit_interval_ms
        self._last_auto_commit = time.time()

        self._offset_manager = self._cluster.get_offset_manager(self._consumer_group)

        owned_partition_partial = functools.partial(
            OwnedPartition, consumer_group=self._consumer_group)
        if partitions:
            self._partitions = {owned_partition_partial(p): p
                                for p in partitions}
        else:
            self._partitions = {owned_partition_partial(p): topic.partitions[k]
                                for k, p in topic.partitions.iteritems()}
        self._partitions_by_id = {p.partition.id: p
                                  for p in self._partitions.iterkeys()}
        # Organize partitions by leader for efficient queries
        self._partitions_by_leader = defaultdict(list)
        for p in self._partitions.iterkeys():
            self._partitions_by_leader[p.partition.leader].append(p)
        self.partition_cycle = itertools.cycle(self._partitions.keys())

        self._running = True

        if self._auto_commit_enable:
            self._autocommit_worker_thread = self._setup_autocommit_worker()
        # we need to get the most up-to-date offsets before starting consumption
        self.fetch_offsets()
        self._fetch_workers = self._setup_fetch_workers()

    @property
    def topic(self):
        return self._topic

    @property
    def partitions(self):
        return self._partitions

    def __del__(self):
        self.stop()

    def stop(self):
        self._running = False

    def _setup_autocommit_worker(self):
        def autocommitter():
            while True:
                if not self._running:
                    break
                if self._auto_commit_enable:
                    self._auto_commit()
        log.debug("Starting autocommitter thread")
        return self._cluster.handler.spawn(autocommitter)

    def _setup_fetch_workers(self):
        def fetcher():
            while True:
                if not self._running:
                    break
                self.fetch()
        return [self._cluster.handler.spawn(fetcher)
                for i in xrange(self._num_consumer_fetchers)]

    def __iter__(self):
        while True:
            message = self.consume()
            if not message and self._consumer_timed_out():
                raise StopIteration
            yield message

    def _consumer_timed_out(self):
        if self._consumer_timeout_ms == -1:
            return False
        disp = (time.time() - self._last_message_time) * 1000.0
        return disp > self._consumer_timeout_ms

    def consume(self):
        """Get one message from the consumer.

        :param timeout: Seconds to wait before returning None
        """
        owned_partition = self.partition_cycle.next()
        message = owned_partition.consume()

        if message:
            self._last_message_time = time.time()
            return message

    def _auto_commit(self):
        if not self._auto_commit_enable or self._auto_commit_interval_ms == 0:
            return

        if (time.time() - self._last_auto_commit) * 1000.0 >= self._auto_commit_interval_ms:
            log.info("Autocommitting consumer offset for consumer group %s and topic %s",
                     self._consumer_group, self._topic.name)
            self.commit_offsets()
            self._last_auto_commit = time.time()

    def commit_offsets(self):
        """Commit offsets for this consumer's topic

        Uses the offset commit/fetch API
        """
        if not self._consumer_group:
            raise Exception("consumer group must be specified to commit offsets")

        self._offset_manager = self._cluster.get_offset_manager(self._consumer_group)

        reqs = [p.build_offset_commit_request() for p in self._partitions.keys()]
        log.info("Committing offsets for %d partitions", len(reqs))
        self._offset_manager.commit_consumer_group_offsets(
            self._consumer_group, 1, 'pykafka', reqs)

    def fetch_offsets(self):
        """Fetch offsets for this consumer's topic

        Uses the offset commit/fetch API
        Should be called when consumer starts and after any errors
        """
        if not self._consumer_group:
            raise Exception("consumer group must be specified to fetch offsets")

        self._offset_manager = self._cluster.get_offset_manager(self._consumer_group)

        log.info("Fetching offsets")

        reqs = [p.build_offset_fetch_request() for p in self._partitions.keys()]
        res = self._offset_manager.fetch_consumer_group_offsets(self._consumer_group, reqs)
        for partition_id, pres in res.topics[self._topic.name].iteritems():
            partition = self._partitions_by_id[partition_id]
            partition.set_offset_counters(pres)

    def fetch(self):
        """Fetch new messages for all partitions

        Create a FetchRequest for each broker and send it. Enqueue each of the
        returned messages in the approprate OwnedPartition.
        """
        for broker, owned_partitions in self._partitions_by_leader.iteritems():
            reqs = []
            for owned_partition in owned_partitions:
                has_room = owned_partition.message_count < self._queued_max_messages
                if owned_partition.empty and has_room:
                    # attempt to acquire lock, just pass if we can't
                    if owned_partition.lock.acquire(False):
                        reqs.append(owned_partition.build_fetch_request(
                            self._fetch_message_max_bytes))
            if reqs:
                response = broker.fetch_messages(
                    reqs, timeout=self._fetch_wait_max_ms,
                    min_bytes=self._fetch_min_bytes
                )
                for partition_id, pres in response.topics[self._topic.name].iteritems():
                    owned_partition = self._partitions_by_id[partition_id]
                    owned_partition.enqueue_messages(pres.messages)
                    owned_partition.lock.release()

class OwnedPartition(object):
    """A partition that is owned by a SimpleConsumer.

    Used to keep track of offsets and the internal message queue.
    """

    def __init__(self,
                 partition,
                 consumer_group=None):
        self.partition = partition
        self._messages = Queue()
        self.last_offset_consumed = 0
        self.next_offset = 0
        self.lock = threading.Lock()

    @property
    def message_count(self):
        return self._messages.qsize()

    @property
    def empty(self):
        return self._messages.empty()

    def set_offset_counters(self, res):
        """Set the internal offset counters from an OffsetFetchResponse

        :param res: an OffsetFetchPartitionResponse containing the committed
            offset for this partition
        :type res: <protocol.PartitionOffsetFetchResponse>
        """
        self.last_offset_consumed = res.offset
        self.next_offset = res.offset + 1
        log.info("Last offset consumed: %d", self.last_offset_consumed)

    def build_fetch_request(self, max_bytes):
        return PartitionFetchRequest(
            self.partition.topic.name, self.partition.id,
            self.next_offset, max_bytes)

    def build_offset_commit_request(self):
        """Create a PartitionOffsetCommitRequest for this partition
        """
        return PartitionOffsetCommitRequest(
            self.partition.topic.name,
            self.partition.id,
            self.last_offset_consumed,
            int(time.time()),
            'pykafka'
        )

    def build_offset_fetch_request(self):
        """Create a PartitionOffsetFetchRequest for this partition
        """
        return PartitionOffsetFetchRequest(
            self.partition.topic.name,
            self.partition.id
        )

    def consume(self):
        """Get a single message from this partition
        """
        try:
            message = self._messages.get_nowait()
            self.last_offset_consumed = message.offset
            return message
        except Empty:
            return None

    def enqueue_messages(self, messages):
        for message in messages:
            if message.offset < self.last_offset_consumed:
                continue
            self._messages.put(message)
            self.next_offset = message.offset + 1
<|endoftext|>


<commit_before>import sqlite3
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import numpy as np
from discord.ext import commands

class Stats:
    def __init__(self, yeebot):
        self.yeebot = yeebot
        self.conn = sqlite3.connect('db/yee.db')
        self.cur = self.conn.cursor()

    @commands.command(pass_context=True, hidden=True)
    async def reqstats(self, ctx):
        channel = ctx.message.channel

        names = []
        times_memed = []

        self.cur.execute("SELECT username, memes_requested FROM users ORDER BY memes_requested DESC LIMIT 5;")

        rows = self.cur.fetchall()
        for row in rows:
            names.append(row[0])
            times_memed.append(row[1])

        yeestring = 'Memes Requested:\n```\n'
        for x in range(0, len(names)):
            yeestring += '{}: {}\n'.format(names[x], times_memed[x])
        yeestring += '```'

        return await self.yeebot.say(yeestring)

    @commands.command(pass_context=True, hidden=True)
    async def substats(self, ctx):
        channel = ctx.message.channel

        names = []
        submissions = []
        rejected = []
        approved = []

        self.cur.execute("select submitter_name, count (submitter_name) from links group "
                    "by submitter_name order by count (submitter_name) desc limit 5;")

        rows = self.cur.fetchall()
        for row in rows:
            names.append(row[0])
            submissions.append(row[1])

        for name in names:
            self.cur.execute("select count (*) from links where status = 'approved' and "
                        "submitter_name = ?;", (name,))
            num_appr = self.cur.fetchone()
            approved.append(int(num_appr[0]))
            self.cur.execute("select count (*) from links where status = 'rejected' and "
                        "submitter_name = ?;", (name,))
            num_rej = self.cur.fetchone()
            rejected.append(int(num_rej[0]))

        N = np.arange(len(names))
        width = 0.4
        approved_bar = plt.bar(N, approved, width, color="green", edgecolor = "none")
        rejected_bar = plt.bar(N, rejected, width, color="red", bottom=approved, edgecolor="none")

        bars = [rejected_bar, approved_bar]

        counter = 0
        bar_counter = 0

        bar_heights = []
        percentages = []

        for x in range(0, len(approved)):
            percentages.append(int(100 * (approved[x] / (approved[x] + rejected[x]))))

        for container in rejected_bar:
            bar_heights.append(container.get_height())
        for container in approved_bar:
            bar_heights[counter] = bar_heights[counter] + container.get_height()
            counter += 1
        for container in approved_bar:
            plt.text(container.get_x() + container.get_width() + 0.02, bar_heights[bar_counter] / 2, str(int(bar_heights[bar_counter])), ha='left', va='center')
            bar_counter += 1

        for bar in bars:
            for container in bar:
                height = container.get_height()
                y = container.get_y()
                width = container.get_width()
                x = container.get_x()
                plt.text(x + (width / 2), y + (height / 2), str(int(height)), ha='center', va='center')

        plt.title('Top 5 Submitters by Volume')
        plt.ylabel('# Memes Submitted')
        plt.xticks(N, names)

        plt.figlegend((approved_bar, rejected_bar),
                      ('Approved', 'Rejected'),
                      'upper right')
        plt.savefig('cogs/output/stats.png')

        yeestring = '```\n'
        for x in range(0, len(names)):
            yeestring += '{}: {}% approval rate.\n'.format(names[x], percentages[x])
        yeestring += '```'

        await self.yeebot.send_file(channel, 'cogs/output/stats.png')
        return await self.yeebot.say(yeestring)

def setup(yeebot):
    yeebot.add_cog(Stats(yeebot))
<commit_msg>Remove graph if it already exists before creating<commit_after>import sqlite3
import os
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import numpy as np
from discord.ext import commands

class Stats:
    def __init__(self, yeebot):
        self.yeebot = yeebot
        self.conn = sqlite3.connect('db/yee.db')
        self.cur = self.conn.cursor()

    @commands.command(pass_context=True, hidden=True)
    async def reqstats(self, ctx):
        channel = ctx.message.channel

        names = []
        times_memed = []

        self.cur.execute("SELECT username, memes_requested FROM users ORDER BY memes_requested DESC LIMIT 5;")

        rows = self.cur.fetchall()
        for row in rows:
            names.append(row[0])
            times_memed.append(row[1])

        yeestring = 'Memes Requested:\n```\n'
        for x in range(0, len(names)):
            yeestring += '{}: {}\n'.format(names[x], times_memed[x])
        yeestring += '```'

        return await self.yeebot.say(yeestring)

    @commands.command(pass_context=True, hidden=True)
    async def substats(self, ctx):
        channel = ctx.message.channel

        names = []
        submissions = []
        rejected = []
        approved = []

        self.cur.execute("select submitter_name, count (submitter_name) from links group "
                    "by submitter_name order by count (submitter_name) desc limit 5;")

        rows = self.cur.fetchall()
        for row in rows:
            names.append(row[0])
            submissions.append(row[1])

        for name in names:
            self.cur.execute("select count (*) from links where status = 'approved' and "
                        "submitter_name = ?;", (name,))
            num_appr = self.cur.fetchone()
            approved.append(int(num_appr[0]))
            self.cur.execute("select count (*) from links where status = 'rejected' and "
                        "submitter_name = ?;", (name,))
            num_rej = self.cur.fetchone()
            rejected.append(int(num_rej[0]))

        N = np.arange(len(names))
        width = 0.4
        approved_bar = plt.bar(N, approved, width, color="green", edgecolor = "none")
        rejected_bar = plt.bar(N, rejected, width, color="red", bottom=approved, edgecolor="none")

        bars = [rejected_bar, approved_bar]

        counter = 0
        bar_counter = 0

        bar_heights = []
        percentages = []

        for x in range(0, len(approved)):
            percentages.append(int(100 * (approved[x] / (approved[x] + rejected[x]))))

        for container in rejected_bar:
            bar_heights.append(container.get_height())
        for container in approved_bar:
            bar_heights[counter] = bar_heights[counter] + container.get_height()
            counter += 1
        for container in approved_bar:
            plt.text(container.get_x() + container.get_width() + 0.02, bar_heights[bar_counter] / 2, str(int(bar_heights[bar_counter])), ha='left', va='center')
            bar_counter += 1

        for bar in bars:
            for container in bar:
                height = container.get_height()
                y = container.get_y()
                width = container.get_width()
                x = container.get_x()
                plt.text(x + (width / 2), y + (height / 2), str(int(height)), ha='center', va='center')

        plt.title('Top 5 Submitters by Volume')
        plt.ylabel('# Memes Submitted')
        plt.xticks(N, names)

        if os.path.isfile('cogs/output/stats.png')
            os.remove('cogs/output/stats.png')

        plt.figlegend((approved_bar, rejected_bar),
                      ('Approved', 'Rejected'),
                      'upper right')
        plt.savefig('cogs/output/stats.png')

        yeestring = '```\n'
        for x in range(0, len(names)):
            yeestring += '{}: {}% approval rate.\n'.format(names[x], percentages[x])
        yeestring += '```'

        await self.yeebot.send_file(channel, 'cogs/output/stats.png')
        return await self.yeebot.say(yeestring)

def setup(yeebot):
    yeebot.add_cog(Stats(yeebot))
<|endoftext|>


<commit_before>#!/usr/bin/env python
#-*- coding:utf-8 -*-

import logging
import threading
import time
import socket
from functools import partial

import rflib.ipc.IPC as IPC
import rflib.ipc.MongoIPC as MongoIPC
from rflib.ipc.RFProtocol import *
from rflib.ipc.RFProtocolFactory import RFProtocolFactory
from rflib.defs import *
from rflib.types.Match import *
from rflib.types.Action import *
from rflib.types.Option import *

class RFMonitor(RFProtocolFactory, IPC.IPCMessageProcessor):
    """Monitors all the controller instances for failiure"""
    def __init__(self, *arg, **kwargs):
        self.controllers = dict()
        self.monitors = dict()
        self.ipc = MongoIPC.MongoIPCMessageService(MONGO_ADDRESS,
                                                   MONGO_DB_NAME,
                                                   RFMONITOR_ID,
                                                   threading.Thread,
                                                   time.sleep)
        self.ipc.listen(RFMONITOR_RFPROXY_CHANNEL, self, self, False)
        self.log = logging.getLogger("rfmonitor")
        self.log.setLevel(logging.INFO)
        ch = logging.StreamHandler()
        ch.setLevel(logging.INFO)
        ch.setFormatter(logging.Formatter(logging.BASIC_FORMAT))
        self.log.addHandler(ch)

    def process(self, _from, to, channel, msg):
        type_ = msg.get_type()
        if type_ == CONTROLLER_REGISTER:
            self.controllers[msg.get_ct_addr() + ':'
                             + str(msg.get_ct_port())] = msg.get_ct_role()
            self.log.info("A %s controller at %s:%s is up", msg.get_ct_role(),
                          msg.get_ct_addr(), msg.get_ct_port())
            test_function = partial(self.test)
            monitor = Monitor(msg.get_ct_addr(), msg.get_ct_port(), test_function,
                              callback_time=1000)
            self.monitors[msg.get_ct_addr() + ':' +
                          str(msg.get_ct_port())] = monitor
            monitor.start_test()

    def test(self, host, port):
        s = socket(AF_INET, SOCK_STREAM)
        result = s.connect_ex((host, port))

        if result != 0:
            self.log.info("Controller listening at %s:%s died", host, port)
            self.controllers.pop(host + ':' + str(port), None)
            self.monitors[host + ':' + str(port)].stop_test()
            self.monitors.pop(host + ':' + str(port), None)
        s.close()

class Monitor(object):
    """Monitors each controller individually"""
    def __init__(self, host, port, callback_time):
        super(Monitor, self).__init__()
        self.host = host
        self.port = port
        self.callback_time = callback_time

    def start_test():
        pass

    def stop_test():
        pass

    def schedule_test():
        pass

    def run_test():
        pass

if __name__ == "__main__":
    description = 'RFMonitor monitors RFProxy instances for failiure'
    epilog = 'Report bugs to: https://github.com/routeflow/RouteFlow/issues'
    RFMonitor()
<commit_msg>Monitor class for scheduling tests on controller.<commit_after>#!/usr/bin/env python
#-*- coding:utf-8 -*-

import logging
import threading
import time
import socket
from functools import partial

import rflib.ipc.IPC as IPC
import rflib.ipc.MongoIPC as MongoIPC
from rflib.ipc.RFProtocol import *
from rflib.ipc.RFProtocolFactory import RFProtocolFactory
from rflib.defs import *
from rflib.types.Match import *
from rflib.types.Action import *
from rflib.types.Option import *

class RFMonitor(RFProtocolFactory, IPC.IPCMessageProcessor):
    """Monitors all the controller instances for failiure"""
    def __init__(self, *arg, **kwargs):
        self.controllers = dict()
        self.monitors = dict()
        self.ipc = MongoIPC.MongoIPCMessageService(MONGO_ADDRESS,
                                                   MONGO_DB_NAME,
                                                   RFMONITOR_ID,
                                                   threading.Thread,
                                                   time.sleep)
        self.ipc.listen(RFMONITOR_RFPROXY_CHANNEL, self, self, False)
        self.log = logging.getLogger("rfmonitor")
        self.log.setLevel(logging.INFO)
        ch = logging.StreamHandler()
        ch.setLevel(logging.INFO)
        ch.setFormatter(logging.Formatter(logging.BASIC_FORMAT))
        self.log.addHandler(ch)

    def process(self, _from, to, channel, msg):
        type_ = msg.get_type()
        if type_ == CONTROLLER_REGISTER:
            self.controllers[msg.get_ct_addr() + ':'
                             + str(msg.get_ct_port())] = msg.get_ct_role()
            self.log.info("A %s controller at %s:%s is up", msg.get_ct_role(),
                          msg.get_ct_addr(), msg.get_ct_port())
            test_function = partial(self.test)
            monitor = Monitor(msg.get_ct_addr(), msg.get_ct_port(), test_function,
                              callback_time=1000)
            self.monitors[msg.get_ct_addr() + ':' +
                          str(msg.get_ct_port())] = monitor
            monitor.start_test()

    def test(self, host, port):
        s = socket(AF_INET, SOCK_STREAM)
        result = s.connect_ex((host, port))

        if result != 0:
            self.log.info("Controller listening at %s:%s died", host, port)
            self.controllers.pop(host + ':' + str(port), None)
            self.monitors[host + ':' + str(port)].stop_test()
            self.monitors.pop(host + ':' + str(port), None)
        s.close()

class Monitor(object):
    """Monitors each controller individually"""
    def __init__(self, host, port, test, callback_time=1000):
        super(Monitor, self).__init__()
        self.host = host
        self.port = port
        self.test = test
        self.callback_time = callback_time
        self.running = False

    def start_test(self):
        self.running = True
        self.timeout = time.time()
        self.schedule_test()

    def stop_test(self):
        self.running = False

    def schedule_test(self):
        while self.running:
            current_time = time.time()
            if self.timeout <= current_time:
                self.timeout += self.callback_time/1000.00
            self.test(self.host, self.port)

if __name__ == "__main__":
    description = 'RFMonitor monitors RFProxy instances for failiure'
    epilog = 'Report bugs to: https://github.com/routeflow/RouteFlow/issues'
    RFMonitor()
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
This is a python script to detect bibs using Keras.
Author: Alex Cummaudo
Date: 1 Aug 2017
"""

from __future__ import division

import pickle
import os
import json

from time import time as now
from optparse import OptionParser
from keras_frcnn import config
from keras_frcnn import roi_helpers
from keras import backend as K
from keras.layers import Input
from keras.models import Model

import cv2
import numpy as np

BOUNDING_BOX_THRESH = 0.8

# Parser
OPTS_PARSER = OptionParser()
OPTS_PARSER.add_option("-i", dest="input_file", help="File to process")
OPTS_PARSER.add_option("-o", dest="output_dir", help="Directory to put output")
OPTS_PARSER.add_option("-j", dest="json_only", help="Output JSON only", default=False)
OPTS_PARSER.add_option("-g", dest="image_only", help="Output image only", default=False)
OPTS_PARSER.add_option("-c", dest="config_file", help="Pickle config file")

def check_options(options):
    """Checks if the provided options are ok.

    Args:
        options (object): The provided options.
    Returns:
        boolean: True if ok, false otherwise.
    """
    if not options.input_file:
        OPTS_PARSER.error("Missing input file")
        return false
    if not options.output_dir:
        OPTS_PARSER.error("Missing output directory")
        return False
    if not options.config_file:
        OPTS_PARSER.error("Missing config file")
        return False
    return True

def load_config(filename):
    """Loads in configuration file.

    Args:
        filename (string): The filename of the config file.
    Returns:
        object: A deserialised pickle object.
    """
    with open(filename, 'rb') as file_ptr:
        return pickle.load(file_ptr)

def format_image(img, config):
    """Formats the image size based on config.

    Args:
        img (3D numpy array): The image to process.
        config (object): Configuration object.

    Returns:
        tuple: Containing the resized image & channels with the new ratio
    """
    # Image resize
    img_min_side = float(config.im_size)
    (height, width, _) = img.shape

    if width <= height:
        ratio = img_min_side / width
        new_height = int(ratio * height)
        new_width = int(img_min_side)
    else:
        ratio = img_min_side / height
        new_width = int(ratio * width)
        new_height = int(img_min_side)

    img = cv2.resize(img, (new_width, new_height),
                     interpolation=cv2.INTER_CUBIC)

    # Channel modifications
    img = img[:, :, (2, 1, 0)]
    img = img.astype(np.float32)
    img[:, :, 0] -= config.img_channel_mean[0]
    img[:, :, 1] -= config.img_channel_mean[1]
    img[:, :, 2] -= config.img_channel_mean[2]
    img /= config.img_scaling_factor
    img = np.transpose(img, (2, 0, 1))
    img = np.expand_dims(img, axis=0)

    return (img, ratio)

def configure_keras_models(config):
    """Configures Keras.

    Args:
        nn (Neural Net): A keras NN.
        config (object): The config file.
    Returns:
        tuple: A tuple of three classifiers: (1) the rpn classifier, (2) the
               classifier (only), (3) the classifier.
    """
    # Import the correct NN according to config
    # This must be called within main so that import is called
    if config.network == 'resnet50':
        import keras_frcnn.resnet as nn
        num_features = 1024
    elif config.network == 'vgg':
        import keras_frcnn.vgg as nn
        num_features = 512

    # Configure Keras
    if K.image_dim_ordering() == 'th':
        input_shape_img = (3, None, None)
        input_shape_features = (num_features, None, None)
    else:
        input_shape_img = (None, None, 3)
        input_shape_features = (None, None, num_features)

    # Config inputs
    img_input = Input(shape=input_shape_img)
    roi_input = Input(shape=(config.num_rois, 4))
    feature_map_input = Input(shape=input_shape_features)

    # define the base network (resnet here, can be VGG, Inception, etc)
    shared_layers = nn.nn_base(img_input, trainable=True)

    # define the RPN, built on the base layers
    num_anchors = len(config.anchor_box_scales) * len(config.anchor_box_ratios)
    rpn_layers = nn.rpn(shared_layers, num_anchors)

    # Classifier to return
    classifier = nn.classifier(
        feature_map_input,
        roi_input,
        config.num_rois,
        nb_classes=len(config.class_mapping),
        trainable=True
    )

    print("Configuring Keras with:")
    print("  - Neural Network: %s..." % config.network)
    print("  - Weights loaded from: %s" % config.model_path)
    print("  - Dimension Ordering: %s" % K.image_dim_ordering())
    print("  - Num Features: %s" % num_features)
    print("  - Num rois: %s" % config.num_rois)

    model_rpn = Model(img_input, rpn_layers)
    model_classifier_only = Model([feature_map_input, roi_input], classifier)
    model_classifier = Model([feature_map_input, roi_input], classifier)

    model_rpn.load_weights(config.model_path, by_name=True)
    model_classifier.load_weights(config.model_path, by_name=True)

    model_rpn.compile(optimizer='sgd', loss='mse')
    model_classifier.compile(optimizer='sgd', loss='mse')

    return (model_rpn, model_classifier_only, model_classifier)

def bboxes_to_dict(bboxes, probs, ratio):
    """Converts the bounding boxes to per-pixel coords as a dictionary.
    Args:
        bboxes (list): A list of all the bounding boxes detected.
        probs (list): A list containing the estimated accuracy.
        ratio (float): The aspect ratio to convert.
    Returns:
        list: A mapped list to per-pixel coordinates as a dictionary.
    """
    result_dets = []
    for key in bboxes:
        bbox = np.array(bboxes[key])
        new_boxes, probs = roi_helpers.non_max_suppression_fast(
            bbox,
            np.array(probs[key]),
            overlap_thresh=0.5
        )
        for jk in range(new_boxes.shape[0]):
            x1, y1, x2, y2 = new_boxes[jk, :]
            real_x1 = int(round(x1 // ratio))
            real_y1 = int(round(y1 // ratio))
            real_x2 = int(round(x2 // ratio))
            real_y2 = int(round(y2 // ratio))
            accuracy = float(probs[jk])
            result_dets.append({
                "x1": real_x1,
                "y1": real_y1,
                "x2": real_x2,
                "y2": real_y2,
                "accuracy": accuracy
            })
    return result_dets

def annotate_image(img, detections):
    for region in detections:
        x1 = region["x1"]
        y1 = region["y1"]
        x2 = region["x2"]
        y2 = region["y2"]
        acc = region["accuracy"]
        lime = (0,255,0)
        black = (0,0,0)
        font = cv2.FONT_HERSHEY_PLAIN
        label = "%s%%" % int(acc * 100)
        cv2.rectangle(img, (x1, y1), (x2, y2), lime, 2)
        cv2.rectangle(img, (x1, y1-5), (x2, y2), lime, -1)
        cv2.putText(img, label, (x1,y1), font, 1, black)
    return img

def process_image(img, config, models):
    """Process the given image using FRCNN.
    Args:
        img (numpy 3D array): The image to process.
        config (object): Configuration settings.
        models (tuple): The models generated from Keras.
    Returns:
        list: Returns a list of detected bounding boxes, mapped to
              their per-pixel coordinates.
    """
    x, ratio = format_image(img, config)
    if K.image_dim_ordering() == 'tf':
        x = np.transpose(x, (0, 2, 3, 1))

    # Unpack the models
    model_rpn, model_classifier_only, model_class = models

    # get the feature maps and output from the RPN
    [y1, y2, f] = model_rpn.predict(x)
    roi = roi_helpers.rpn_to_roi(
        y1,
        y2,
        config,
        K.image_dim_ordering(),
        overlap_thresh=0.7
    )

    # convert from (x1,y1,x2,y2) to (x,y,w,h)
    roi[:, 2] -= roi[:, 0]
    roi[:, 3] -= roi[:, 1]

    # Convert classes to a dict
    class_mapping = { v: k for k, v in config.class_mapping.items() }

    # apply the spatial pyramid pooling to the proposed regions
    bboxes = {}
    probs = {}

    # FRCNN algo
    for jk in range(roi.shape[0] // config.num_rois + 1):
        rois = np.expand_dims(
            roi[config.num_rois * jk:config.num_rois * (jk + 1), :],
            axis=0
        )
        if rois.shape[1] == 0:
            break

        if jk == roi.shape[0] // config.num_rois:
            curr_shape = rois.shape
            target_shape = (curr_shape[0], config.num_rois, curr_shape[2])
            rois_padded = np.zeros(target_shape).astype(rois.dtype)
            rois_padded[:, :curr_shape[1], :] = rois
            rois_padded[0, curr_shape[1]:, :] = rois[0, 0, :]
            rois = rois_padded

        [p_cls, p_regr] = model_classifier_only.predict([f, rois])

        for ii in range(p_cls.shape[1]):
            if np.max(p_cls[0, ii, :]) < BOUNDING_BOX_THRESH or np.argmax(
                    p_cls[0, ii, :]) == (p_cls.shape[2] - 1):
                continue

            cls_name = class_mapping[np.argmax(p_cls[0, ii, :])]

            if cls_name not in bboxes:
                bboxes[cls_name] = []
                probs[cls_name] = []

            (x, y, w, h) = rois[0, ii, :]

            cls_num = np.argmax(p_cls[0, ii, :])
            try:
                (tx, ty, tw, th) = p_regr[0, ii, 4 * cls_num:4 * (cls_num + 1)]
                tx /= config.classifier_regr_std[0]
                ty /= config.classifier_regr_std[1]
                tw /= config.classifier_regr_std[2]
                th /= config.classifier_regr_std[3]
                x, y, w, h = roi_helpers.apply_regr(x, y, w, h, tx, ty, tw, th)
            except:
                pass
            bboxes[cls_name].append([
                config.rpn_stride * x,
                config.rpn_stride * y,
                config.rpn_stride * (x + w),
                config.rpn_stride * (y + h)
            ])
            probs[cls_name].append(np.max(p_cls[0, ii, :]))

    # Convert all bboxes to list of dict items
    return bboxes_to_dict(bboxes, probs, ratio)

def main():
    """Main program entry point"""
    options = OPTS_PARSER.parse_args()[0]

    if not check_options(options):
        return

    # Load config
    config = load_config(options.config_file)

    if not os.path.exists(options.output_dir):
        os.makedirs(options.output_dir)

    img = cv2.imread(options.input_file)

    print("Processing: %s..." % options.input_file)

    if not os.path.exists(options.input_file):
        print("No such image at %s. Aborting." % options.input_file)
        return

    start_time = now()
    models = configure_keras_models(config)
    detections = process_image(img, config, models)

    if detections != None:
        for region in detections:
            x1 = region["x1"]
            y1 = region["y1"]
            x2 = region["x2"]
            y2 = region["y2"]
            acc = region["accuracy"]
            w = x2 - x1
            h = y2 - y1
            print("!! Detection at (%s, %s) to (%s, %s) "
            "(w=%s, h=%s) "
            "[Accuracy = %s%%]" % (x1, y1, x2, y2, w, h, int(acc * 100)))

    elapsed_time = now() - start_time
    print("Time taken: %ss." % elapsed_time)

    input_file_basename = os.path.basename(options.input_file)

    # Writing out annotated image done if not doing json only
    if not options.json_only:
        print("Annotating image...")
        img_annotated = annotate_image(img, detections)
        img_annotated_file = "%s/%s" % (options.output_dir, input_file_basename)
        print("Writing annotated image to '%s'..." % img_annotated_file)
        cv2.imwrite(img_annotated_file, img_annotated)

    # JSON processing done if not doing image only
    if not options.image_only:
        data = {"bib_regions": detections, "elapsed_time": elapsed_time}
        json_filename = "%s/%s.json" % (
            options.output_dir,
            os.path.splitext(input_file_basename)[0]
        )
        print("Writing JSON file to '%s'..." % json_filename)
        with open(json_filename, 'w') as outfile:
             json.dump(data, outfile)

# Start of script
if __name__ == '__main__':
    main()

<commit_msg>FIX: Fix bounding box covering bib<commit_after>#!/usr/bin/env python
"""
This is a python script to detect bibs using Keras.
Author: Alex Cummaudo
Date: 1 Aug 2017
"""

from __future__ import division

import pickle
import os
import json

from time import time as now
from optparse import OptionParser
from keras_frcnn import config
from keras_frcnn import roi_helpers
from keras import backend as K
from keras.layers import Input
from keras.models import Model

import cv2
import numpy as np

BOUNDING_BOX_THRESH = 0.8

# Parser
OPTS_PARSER = OptionParser()
OPTS_PARSER.add_option("-i", dest="input_file", help="File to process")
OPTS_PARSER.add_option("-o", dest="output_dir", help="Directory to put output")
OPTS_PARSER.add_option("-j", dest="json_only", help="Output JSON only", default=False)
OPTS_PARSER.add_option("-g", dest="image_only", help="Output image only", default=False)
OPTS_PARSER.add_option("-c", dest="config_file", help="Pickle config file")

def check_options(options):
    """Checks if the provided options are ok.

    Args:
        options (object): The provided options.
    Returns:
        boolean: True if ok, false otherwise.
    """
    if not options.input_file:
        OPTS_PARSER.error("Missing input file")
        return false
    if not options.output_dir:
        OPTS_PARSER.error("Missing output directory")
        return False
    if not options.config_file:
        OPTS_PARSER.error("Missing config file")
        return False
    return True

def load_config(filename):
    """Loads in configuration file.

    Args:
        filename (string): The filename of the config file.
    Returns:
        object: A deserialised pickle object.
    """
    with open(filename, 'rb') as file_ptr:
        return pickle.load(file_ptr)

def format_image(img, config):
    """Formats the image size based on config.

    Args:
        img (3D numpy array): The image to process.
        config (object): Configuration object.

    Returns:
        tuple: Containing the resized image & channels with the new ratio
    """
    # Image resize
    img_min_side = float(config.im_size)
    (height, width, _) = img.shape

    if width <= height:
        ratio = img_min_side / width
        new_height = int(ratio * height)
        new_width = int(img_min_side)
    else:
        ratio = img_min_side / height
        new_width = int(ratio * width)
        new_height = int(img_min_side)

    img = cv2.resize(img, (new_width, new_height),
                     interpolation=cv2.INTER_CUBIC)

    # Channel modifications
    img = img[:, :, (2, 1, 0)]
    img = img.astype(np.float32)
    img[:, :, 0] -= config.img_channel_mean[0]
    img[:, :, 1] -= config.img_channel_mean[1]
    img[:, :, 2] -= config.img_channel_mean[2]
    img /= config.img_scaling_factor
    img = np.transpose(img, (2, 0, 1))
    img = np.expand_dims(img, axis=0)

    return (img, ratio)

def configure_keras_models(config):
    """Configures Keras.

    Args:
        nn (Neural Net): A keras NN.
        config (object): The config file.
    Returns:
        tuple: A tuple of three classifiers: (1) the rpn classifier, (2) the
               classifier (only), (3) the classifier.
    """
    # Import the correct NN according to config
    # This must be called within main so that import is called
    if config.network == 'resnet50':
        import keras_frcnn.resnet as nn
        num_features = 1024
    elif config.network == 'vgg':
        import keras_frcnn.vgg as nn
        num_features = 512

    # Configure Keras
    if K.image_dim_ordering() == 'th':
        input_shape_img = (3, None, None)
        input_shape_features = (num_features, None, None)
    else:
        input_shape_img = (None, None, 3)
        input_shape_features = (None, None, num_features)

    # Config inputs
    img_input = Input(shape=input_shape_img)
    roi_input = Input(shape=(config.num_rois, 4))
    feature_map_input = Input(shape=input_shape_features)

    # define the base network (resnet here, can be VGG, Inception, etc)
    shared_layers = nn.nn_base(img_input, trainable=True)

    # define the RPN, built on the base layers
    num_anchors = len(config.anchor_box_scales) * len(config.anchor_box_ratios)
    rpn_layers = nn.rpn(shared_layers, num_anchors)

    # Classifier to return
    classifier = nn.classifier(
        feature_map_input,
        roi_input,
        config.num_rois,
        nb_classes=len(config.class_mapping),
        trainable=True
    )

    print("Configuring Keras with:")
    print("  - Neural Network: %s..." % config.network)
    print("  - Weights loaded from: %s" % config.model_path)
    print("  - Dimension Ordering: %s" % K.image_dim_ordering())
    print("  - Num Features: %s" % num_features)
    print("  - Num rois: %s" % config.num_rois)

    model_rpn = Model(img_input, rpn_layers)
    model_classifier_only = Model([feature_map_input, roi_input], classifier)
    model_classifier = Model([feature_map_input, roi_input], classifier)

    model_rpn.load_weights(config.model_path, by_name=True)
    model_classifier.load_weights(config.model_path, by_name=True)

    model_rpn.compile(optimizer='sgd', loss='mse')
    model_classifier.compile(optimizer='sgd', loss='mse')

    return (model_rpn, model_classifier_only, model_classifier)

def bboxes_to_dict(bboxes, probs, ratio):
    """Converts the bounding boxes to per-pixel coords as a dictionary.
    Args:
        bboxes (list): A list of all the bounding boxes detected.
        probs (list): A list containing the estimated accuracy.
        ratio (float): The aspect ratio to convert.
    Returns:
        list: A mapped list to per-pixel coordinates as a dictionary.
    """
    result_dets = []
    for key in bboxes:
        bbox = np.array(bboxes[key])
        new_boxes, probs = roi_helpers.non_max_suppression_fast(
            bbox,
            np.array(probs[key]),
            overlap_thresh=0.5
        )
        for jk in range(new_boxes.shape[0]):
            x1, y1, x2, y2 = new_boxes[jk, :]
            real_x1 = int(round(x1 // ratio))
            real_y1 = int(round(y1 // ratio))
            real_x2 = int(round(x2 // ratio))
            real_y2 = int(round(y2 // ratio))
            accuracy = float(probs[jk])
            result_dets.append({
                "x1": real_x1,
                "y1": real_y1,
                "x2": real_x2,
                "y2": real_y2,
                "accuracy": accuracy
            })
    return result_dets

def annotate_image(img, detections):
    for region in detections:
        x1 = region["x1"]
        y1 = region["y1"]
        x2 = region["x2"]
        y2 = region["y2"]
        acc = region["accuracy"]
        lime = (0,255,0)
        black = (0,0,0)
        font = cv2.FONT_HERSHEY_PLAIN
        label = "%s%%" % int(acc * 100)
        fnt_sz, baseline = cv2.getTextSize(label, font, 1)
        acc_rect_pt1 = (x1 - 5, y1 + baseline - 5)
        acc_rect_pt2 = (x1 + fnt_sz[0] + 5, y1 - fnt_sz[1] - 5)
        cv2.rectangle(img, (x1, y1), (x2, y2), lime, 2)
        cv2.rectangle(img, acc_rect_pt1, acc_rect_pt2, lime, -1)
        cv2.putText(img, label, (x1,y1), font, 1, black)
    return img

def process_image(img, config, models):
    """Process the given image using FRCNN.
    Args:
        img (numpy 3D array): The image to process.
        config (object): Configuration settings.
        models (tuple): The models generated from Keras.
    Returns:
        list: Returns a list of detected bounding boxes, mapped to
              their per-pixel coordinates.
    """
    x, ratio = format_image(img, config)
    if K.image_dim_ordering() == 'tf':
        x = np.transpose(x, (0, 2, 3, 1))

    # Unpack the models
    model_rpn, model_classifier_only, model_class = models

    # get the feature maps and output from the RPN
    [y1, y2, f] = model_rpn.predict(x)
    roi = roi_helpers.rpn_to_roi(
        y1,
        y2,
        config,
        K.image_dim_ordering(),
        overlap_thresh=0.7
    )

    # convert from (x1,y1,x2,y2) to (x,y,w,h)
    roi[:, 2] -= roi[:, 0]
    roi[:, 3] -= roi[:, 1]

    # Convert classes to a dict
    class_mapping = { v: k for k, v in config.class_mapping.items() }

    # apply the spatial pyramid pooling to the proposed regions
    bboxes = {}
    probs = {}

    # FRCNN algo
    for jk in range(roi.shape[0] // config.num_rois + 1):
        rois = np.expand_dims(
            roi[config.num_rois * jk:config.num_rois * (jk + 1), :],
            axis=0
        )
        if rois.shape[1] == 0:
            break

        if jk == roi.shape[0] // config.num_rois:
            curr_shape = rois.shape
            target_shape = (curr_shape[0], config.num_rois, curr_shape[2])
            rois_padded = np.zeros(target_shape).astype(rois.dtype)
            rois_padded[:, :curr_shape[1], :] = rois
            rois_padded[0, curr_shape[1]:, :] = rois[0, 0, :]
            rois = rois_padded

        [p_cls, p_regr] = model_classifier_only.predict([f, rois])

        for ii in range(p_cls.shape[1]):
            if np.max(p_cls[0, ii, :]) < BOUNDING_BOX_THRESH or np.argmax(
                    p_cls[0, ii, :]) == (p_cls.shape[2] - 1):
                continue

            cls_name = class_mapping[np.argmax(p_cls[0, ii, :])]

            if cls_name not in bboxes:
                bboxes[cls_name] = []
                probs[cls_name] = []

            (x, y, w, h) = rois[0, ii, :]

            cls_num = np.argmax(p_cls[0, ii, :])
            try:
                (tx, ty, tw, th) = p_regr[0, ii, 4 * cls_num:4 * (cls_num + 1)]
                tx /= config.classifier_regr_std[0]
                ty /= config.classifier_regr_std[1]
                tw /= config.classifier_regr_std[2]
                th /= config.classifier_regr_std[3]
                x, y, w, h = roi_helpers.apply_regr(x, y, w, h, tx, ty, tw, th)
            except:
                pass
            bboxes[cls_name].append([
                config.rpn_stride * x,
                config.rpn_stride * y,
                config.rpn_stride * (x + w),
                config.rpn_stride * (y + h)
            ])
            probs[cls_name].append(np.max(p_cls[0, ii, :]))

    # Convert all bboxes to list of dict items
    return bboxes_to_dict(bboxes, probs, ratio)

def main():
    """Main program entry point"""
    options = OPTS_PARSER.parse_args()[0]

    if not check_options(options):
        return

    # Load config
    config = load_config(options.config_file)

    if not os.path.exists(options.output_dir):
        os.makedirs(options.output_dir)

    img = cv2.imread(options.input_file)

    print("Processing: %s..." % options.input_file)

    if not os.path.exists(options.input_file):
        print("No such image at %s. Aborting." % options.input_file)
        return

    start_time = now()
    models = configure_keras_models(config)
    detections = process_image(img, config, models)

    if detections != None:
        for region in detections:
            x1 = region["x1"]
            y1 = region["y1"]
            x2 = region["x2"]
            y2 = region["y2"]
            acc = region["accuracy"]
            w = x2 - x1
            h = y2 - y1
            print("!! Detection at (%s, %s) to (%s, %s) "
            "(w=%s, h=%s) "
            "[Accuracy = %s%%]" % (x1, y1, x2, y2, w, h, int(acc * 100)))

    elapsed_time = now() - start_time
    print("Time taken: %ss." % elapsed_time)

    input_file_basename = os.path.basename(options.input_file)

    # Writing out annotated image done if not doing json only
    if not options.json_only:
        print("Annotating image...")
        img_annotated = annotate_image(img, detections)
        img_annotated_file = "%s/%s" % (options.output_dir, input_file_basename)
        print("Writing annotated image to '%s'..." % img_annotated_file)
        cv2.imwrite(img_annotated_file, img_annotated)

    # JSON processing done if not doing image only
    if not options.image_only:
        data = {"bib_regions": detections, "elapsed_time": elapsed_time}
        json_filename = "%s/%s.json" % (
            options.output_dir,
            os.path.splitext(input_file_basename)[0]
        )
        print("Writing JSON file to '%s'..." % json_filename)
        with open(json_filename, 'w') as outfile:
             json.dump(data, outfile)

# Start of script
if __name__ == '__main__':
    main()

<|endoftext|>