In [1]:
import numpy as np
import pandas as pd
import cv2    
from scipy.misc import imread
import csv
from matplotlib import pyplot as plt

In [2]:
#Name of the rows
names = ["32","60","61","66","67","68","69","73","77",
         "80","87","94","98","99","101","102","103","103-1","105","106","107",
         "108","109","110","111","111-1","112","115","116","116-1","117","117-1",
         "118","121","122","124","125","126","127","128","200",
         "201","202","203","204","205","206","207","208","210","211","212","213",
         "214","215","216","217","218","219","222","223","224","225","226",
         "227","228","229","230","231","232","233","234","235",
         "240","241","242","244","245","246","250","251","252","253","254","255",
         "256","257","259","301","302","303","304","305","306","307","308",
         "309","310","311","312","313","314"]

In [3]:
def parse_image(img, number_rows):

    # Start with all zeros in a number_rows x 16 array
    data = np.zeros((number_rows, 16), dtype=int)
    

    # List holding min and max RGB values for each category

    categories = [
        (5, (250, 0, 0), (255, 5, 5)),             # Red
        (4, (152, 60, 5), (160, 78, 12)),          # Dark brown
        (3, (225, 105, 9), (240, 120, 18)),        # Light brown
        (2, (250, 235, 220), (255, 245, 230)),     # Light orange
        (1, (250, 250, 250), (255, 255, 255)),     # White
        ]

    for category, bgr_min, bgr_max in categories:
        # Extract pixels in the required range and convert them to 255
        mask = cv2.inRange(img, bgr_min, bgr_max)
        image_cat = cv2.bitwise_or(img, np.full(img.shape, 255, dtype=np.uint8), mask=mask)
        
        # Convert the image into greyscale
        image_grey = cv2.cvtColor(image_cat, cv2.COLOR_BGR2GRAY)
        
        # Resize the image to 16 x number_rows
        values = cv2.resize(image_grey, (16, number_rows))

        # Convert non black values into the current category value
        values[values > 0] = category

        # Add the values to the data array
        data = data + values

    return data

In [4]:
def create_categories(file_name, size):
    # Load the heatmap    
    img_src = imread(file_name)
    print(img_src.shape)
    # Crop into 3 sub images
    starty = 28
    cropx = 234
    cropy = 184
    #cropx = 233
    cropy = 300
    images = []
    images.append(img_src[26:,30:30+234])
    images.append(img_src[26:,303:303+234])
    images.append(img_src[26:,572:572+234])

    abc = [parse_image(img, size) for img in images]
    return abc,images

In [5]:
all_, all_images  = create_categories("all.jpg" , 102)


(1230, 813, 3)

In [6]:
def detect_category(input_image):
    categories = [
        (5, (250, 0, 0), (255, 5, 5)),             # Red
        (4, (152, 60, 5), (160, 78, 12)),          # Dark brown
        (3, (225, 105, 9), (240, 120, 18)),        # Light brown
        (2, (250, 235, 220), (255, 245, 230)),     # Light orange
        (1, (250, 250, 250), (255, 255, 255)),     # White
        ]
    count = [0] * 6
    for category, bgr_min, bgr_max in categories:
        mask = cv2.inRange(input_image, bgr_min, bgr_max)
        image_cat = cv2.bitwise_or(input_image, np.full(input_image.shape, 255, dtype=np.uint8), mask=mask)
        image_grey = cv2.cvtColor(image_cat, cv2.COLOR_BGR2GRAY)
        count[category] =  np.count_nonzero(image_grey)
    return np.argmax(count)

In [7]:
data = np.zeros((3,102, 16), dtype=int) 

for z in range(3):
    for i in range(102):
        offset =  -1 * (i/3)
        y_begin = offset + (i*12) + 1
        y_end = offset + (1+i)*12 - 1
        for j in range(16):
            data[z,i,j] = detect_category(all_images[z][y_begin:y_end, j*15 : (j+1)*15])

In [8]:
#Assigning names to the rows and columns of our data frame
I = pd.Index(names, name="rows")
C = pd.Index([chr(i) for i in range(65,81)], name="columns")

In [9]:
#Write the output to CSV files
df = pd.DataFrame(data=data[0,:,:], index=I, columns=C)
df.to_csv('output_A.csv')
df = pd.DataFrame(data=data[1,:,:], index=I, columns=C)
df.to_csv('output_B.csv')
df = pd.DataFrame(data=data[2,:,:], index=I, columns=C)
df.to_csv('output_C.csv')

In [ ]: