Initial commit

src/tools.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from os import listdir
 
 import numpy as np
 import pandas as pd
@@ -158,9 +159,15 @@ def get_image_by_sku(sku: str or int) -> np.array:
 	"""
 	image_data = get_image_data()
 	try:
-		image_name = image_data[image_data['sku'] == str(sku)]['image'].iloc[0]
-		img = Image.open(f'data/images/{image_name}')
-		return np.array(img)
+		for num, dir in enumerate(listdir('data/images')):
+			try:
+				image_name = image_data[image_data['sku'] == str(sku)]['image'].iloc[0]
+				img = Image.open(f'data/images/{image_name}')
+				return np.array(img)
+			except:
+				if num == 2:
+					raise Exception
+
 	except:
 		img = Image.open('data/service_images/' + 'no_img.jpg')
 		return np.array(img)