From d5a4b46c7a71213a7576f7ca843e91f4686b2dc1 Mon Sep 17 00:00:00 2001
From: Rusty Striker <avivr903@tutanota.com>
Date: Fri, 19 Jan 2024 15:51:35 +0200
Subject: [PATCH] implement avg color function

---
 classify.py | 50 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/classify.py b/classify.py
index 7effe41..3afa902 100644
--- a/classify.py
+++ b/classify.py
@@ -2,8 +2,22 @@ import cv2 as cv
 import numpy as np
 import h5py as h5
 
+db = None
+
+def init(path):
+    ''' initializes the database, must be called before any use '''
+    global db
+    db = h5.File(path, 'r')
+
+def needs_init():
+    ''' checks if the database has been initialized '''
+    if db is None:
+        print('db is none, please use init(path_to_db) first!')
+    return db is None
+
 # Extract letter from a bounding box
 def extract_bb(img, bb):
+    ''' extracts a bounding box/letter from the given image '''
     # Get the bounding box
     rect = cv.minAreaRect(bb.astype(np.float32).transpose())
     # will be useful later, map center and size to ints
@@ -16,3 +30,39 @@ def extract_bb(img, bb):
     cropped = cv.getRectSubPix(rot_img, size, center)
     return cropped
 
+def get_img(index):
+    ''' gets image from database '''
+    if needs_init():
+        return None
+    names = list(db['data'].keys())
+    im = names[index]
+    return db['data'][im][:]
+
+def get_attrs(index):
+    ''' gets attribute dict from the database '''
+    if needs_init():
+        return None
+    names = list(db['data'].keys())
+    im = names[index]
+    return db['data'][im].attrs
+
+def get_avg_color(img, mask):
+    '''
+    gets avg color from an image that is underneath a mask,
+    img and mask needs to be of same size(in x,y) but img can
+    have any third dimension size it want(usually 3 for rgb or 1 for grayscale)
+    '''
+    sx, sy, sw = img.shape
+    mx, my = mask.shape
+    if sx != mx or sy != my:
+        print('Image and mask shape doesnt match!')
+        return None
+    avg = np.array(sw, dtype=np.float32)
+    count = 0.0
+    for x in range(sx):
+        for y in range(sy):
+            m = mas[x, y]
+            avg += img[x, y] * m
+            count += m
+    avg /= count
+    return avg