I have datasets of images of gemstones.I have categorised the gemstones into folders based on their colours.(That is Red, Blue, Pink,Purple,Yellow).
What I want to do is:
I want to train a model using hsv model and opencv to detect the colour of the gemstone.That is whether its blue,purple,pink,yellow or Red and any other colour except to these 5 colours to define as undefined colour
Source Code:(referred https://www.kaggle.com)
import os
import matplotlib.pyplot as plt
import seaborn as sn
import cv2
from random import randint
import numpy as np
CLASSES, gems = [], [] # names of classes, count of images for each class
for root, dirs, files in os.walk('C:/Users/User/Desktop/Research Project/images'):
f = os.path.basename(root) # get class name - Red,Blue etc
if len(files) > 0:
gems.append(len(files))
if f not in CLASSES:
CLASSES.append(f) # add folder name
gems_count = len(CLASSES) # 6 = number of classes
print('{} classes with {} images in total'.format(len(CLASSES), sum(gems)))
img_w, img_h = 220, 220 # width and height of image
train_dir = 'C:/Users/User/Desktop/Gem/images/train'
def read_imgs_lbls(_dir):
Images, Labels = [], []
for root, dirs, files in os.walk(_dir):
f = os.path.basename(root) # get class name - Red, Blue, etc
for file in files:
Labels.append(f)
try:
image = cv2.imread(root '/' file) # read the image (OpenCV)
image = cv2.resize(image,(int(img_w*1.5), int(img_h*1.5))) # resize the image (images are different sizes)
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) # converts an image from BGR color space to HSV
Images.append(image)
except Exception as e:
print(e)
Images = np.array(Images)
return (Images, Labels)
def get_class_index(Labels):
for i, n in enumerate(Labels):
for j, k in enumerate(CLASSES): # foreach CLASSES
if n == k:
Labels[i] = j
Labels = np.array(Labels)
return Labels
Train_Imgs, Train_Lbls = read_imgs_lbls(train_dir)
Train_Lbls = get_class_index(Train_Lbls)
print('Shape of train images: {}'.format(Train_Imgs.shape))
print('Shape of train labels: {}'.format(Train_Lbls.shape))
dim = 4
f,ax = plt.subplots(dim,dim)
f.subplots_adjust(0,0,2,2)
for i in range(0,dim):
for j in range(0,dim):
rnd_number = randint(0,len(Train_Imgs))
cl = Train_Lbls[rnd_number]
ax[i,j].imshow(Train_Imgs[rnd_number])
ax[i,j].set_title(CLASSES[cl] ': ' str(cl))
ax[i,j].axis('off')
It read the values from the folder names.But I want to add the lower and upper values of each colour to the training model as in the link below.(Referred how to know if a color is detected on opencv)
import cv2
import numpy as np
img = cv2.imread("img.jpg")
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
lower_val = np.array([37,42,0])
upper_val = np.array([84,255,255])
# Threshold the HSV image - any green color will show up as white
mask = cv2.inRange(hsv, lower_val, upper_val)
# if there are any white pixels on mask, sum will be > 0
hasGreen = np.sum(mask)
if hasGreen > 0:
print('Green detected!')
I found the lower and upper limit hsv values of the colours I want as well
'red': [[9, 255, 255], [0, 50, 70]],
'blue': [[128, 255, 255], [90, 50, 70]],
'yellow': [[35, 255, 255], [25, 50, 70]],
'purple': [[158, 255, 255], [129, 50, 70]]
can anyone please let me know how can i combine detecting the colour using the hsv values(as in the how to know if a color is detected on opencv ) to the my source code.
I am new to image processing and any help is appreciated.
Thankyou.
CodePudding user response:
This code demonstrates how to walk through all files in folder ./images
and return the detected colours:
import os
import numpy as np
import cv2
# map colour names to HSV ranges
color_list = [
['red', [0, 160, 70], [10, 250, 250]],
['pink', [0, 50, 70], [10, 160, 250]],
['yellow', [15, 50, 70], [30, 250, 250]],
['green', [40, 50, 70], [70, 250, 250]],
['cyan', [80, 50, 70], [90, 250, 250]],
['blue', [100, 50, 70], [130, 250, 250]],
['purple', [140, 50, 70], [160, 250, 250]],
['red', [170, 160, 70], [180, 250, 250]],
['pink', [170, 50, 70], [180, 160, 250]]
]
def detect_main_color(hsv_image, colors):
color_found = 'undefined'
max_count = 0
for color_name, lower_val, upper_val in colors:
# threshold the HSV image - any matching color will show up as white
mask = cv2.inRange(hsv_image, np.array(lower_val), np.array(upper_val))
# count white pixels on mask
count = np.sum(mask)
if count > max_count:
color_found = color_name
max_count = count
return color_found
for root, dirs, files in os.walk('./images'):
f = os.path.basename(root)
for file in files:
img = cv2.imread(os.path.join(root, file))
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
print(f"{file}: {detect_main_color(hsv, color_list)}")
Output with three sample images in subfolder images
:
ruby_3.jpg: red
sapphire blue_18.jpg: blue
sapphire pink_18.jpg: pink
sapphire purple_28.jpg: purple
sapphire yellow_9.jpg: yellow
Credits: