parse-scores.py

"""Use image analysis to extract scores from coffee charts."""

from PIL import Image, ImageFilter, ImageEnhance
from pytesseract import image_to_string
import cv2
import os
import sys
import numpy as np


def get_scores(image):
    """Extract the score information from the spider graph."""
    scores = dict()
    scores['dry_fragrence'] = image[80:115, 400:442]
    scores['wet_aroma'] = image[133:153, 585:622]
    scores['brightness'] = image[252:270, 680:715]
    scores['flavor'] = image[415:435, 635:668]
    scores['body'] = image[539:561, 537:578]
    scores['finish'] = image[576:601, 363:403]
    scores['sweetness'] = image[549:571, 205:244]
    scores['clean_cup'] = image[429:452, 116:158]
    scores['complexity'] = image[429:452, 116:158]
    scores['uniformity'] = image[129:152, 220:259]
    scores['cupper_correction'] = image[630:657, 431:476]
    scores['final'] = image[661:700, 437:520]
    return scores


def save_crops(scores):
    names = list()
    for key, value in scores.iteritems():
        cv2.imwrite("tmp/{0}.jpg".format(key), value)
        names.append("tmp/{0}.jpg".format(key))
    return names


def main():
    """Do your thing."""
    fname = sys.argv[1]
    path = os.path.dirname(os.path.abspath(__file__))
    full_path = '/'.join([path, fname])

    image = cv2.imread(full_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    xc, yc, r = 360, 350, 220
    H, W = image.shape
    x, y = np.meshgrid(np.arange(W), np.arange(H))
    d2 = (x - xc)**2 + (y - yc)**2
    mask = d2 < r**2

    outside = np.ma.masked_where(mask, image)
    average_color = outside.mean()
    image[mask] = average_color

    cv2.imwrite("score-sample/processed.jpg", image)

    full_path = '/'.join([path, "score-sample/processed.jpg"])
    image = Image.open(full_path)
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(2)
    image.save("score-sample/processed.jpg")
    full_path = '/'.join([path, "score-sample/processed.jpg"])
    image = cv2.imread(full_path)
    # cv2.imshow("Output", image)
    # cv2.waitKey(0)

    scores = get_scores(image)
    fnames = save_crops(scores)
    for name in fnames:
        full_path = '/'.join([path, name])
        image = Image.open(full_path)
        result = image_to_string(image)
        print name, result


if __name__ == "__main__":
    main()