colorClassification.py 2.19 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
GEPARD - Gepard-Enabled PARticle Detection
Copyright (C) 2018  Lars Bittrich and Josef Brandt, Leibniz-Institut für 
Polymerforschung Dresden e. V. <bittrich-lars@ipfdd.de>    

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program, see COPYING.  
If not, see <https://www.gnu.org/licenses/>.
"""

class ColorRangeHSV(object):
    def __init__(self, name, hue, hue_tolerance, min_sat, max_sat):
        self.name = name
        self.minHue = hue-hue_tolerance/2
        self.maxHue = hue+hue_tolerance/2
        self.minSat = min_sat
        self.maxSat = max_sat
    
    def containsHSV(self, hsv):
        hue = hsv[0]
        sat = hsv[1]
        
        if self.minHue <= hue <= self.maxHue and self.minSat <= sat <= self.maxSat:
            return True
        else:
            if self.name != 'white':
                return False
            else:
                if sat < 128 and hsv[2] > 70:
                    return True

class ColorClassifier(object):
    def __init__(self):
        hue_tolerance = 50
        self.colors = [ColorRangeHSV('yellow', 30, hue_tolerance, 40, 255),
                  ColorRangeHSV('blue', 120, hue_tolerance, 40, 255),
                  ColorRangeHSV('red', 180, hue_tolerance, 40, 255),
                  ColorRangeHSV('red', 0, hue_tolerance, 40, 255),
                  ColorRangeHSV('green', 70, hue_tolerance, 40, 255),
                  ColorRangeHSV('white', 128, 256, 0, 40)]
    
    def classifyColor(self, meanHSV):
        result = 'non-determinable'
        for color in self.colors:
            if color.containsHSV(meanHSV):
                result = color.name
                break
        
        return result