classifier.py

# This file is part of AI Organizer Activity.
# Copyright (C) 2025 Bishoy Wadea
#
# 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; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

from unittest import result
import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input, decode_predictions
from tensorflow import keras
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import os

class SVHNNumberClassifier:
    def __init__(self, model_path='svhn_model.h5'):
        self.model_path = model_path
        self.model = None
        
    def load_model(self):
        if os.path.exists(self.model_path):
            try:
                self.model = keras.models.load_model(self.model_path)
                return True
            except Exception as e:
                print(f"[ERROR] Failed to load model: {e}")
                return False
        else:
            print(f"[ERROR] Model file not found at: {self.model_path}")
            return False
    
    def predict_digit(self, image):
        if self.model is None:
            if not self.load_model():
                return None, 0
        
        if len(image.shape) == 2:
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        
        image_resized = cv2.resize(image, (32, 32))
        image_norm = image_resized.astype('float32') / 255.0
        image_batch = np.expand_dims(image_norm, axis=0)
        
        predictions = self.model.predict(image_batch, verbose=0)[0]
        digit = np.argmax(predictions)
        confidence = predictions[digit]
        
        return int(digit), float(confidence)

class HybridClassifier:
    def __init__(self, use_svhn=True):
        self.general_model = None
        self.general_model_loaded = False
        
        try:
            approaches = [
                lambda: MobileNetV2(weights='imagenet', input_shape=(224, 224, 3)),
                lambda: MobileNetV2(weights='imagenet'),
                lambda: MobileNetV2(weights='imagenet', input_tensor=tf.keras.Input(shape=(224, 224, 3)))
            ]
            
            for i, approach in enumerate(approaches):
                try:
                    self.general_model = approach()
                    self.general_model_loaded = True
                    break
                except Exception as e:
                    print(f"Approach {i+1} failed: {e}")
                    continue
            
            if not self.general_model_loaded:
                print("All MobileNetV2 approaches failed, general classification disabled")
                
        except Exception as e:
            print(f"Critical error loading MobileNetV2: {e}")
            self.general_model = None
            self.general_model_loaded = False
        
        self.use_svhn = use_svhn
        if use_svhn:
            self.svhn_model = SVHNNumberClassifier()
            if not self.svhn_model.load_model():
                print("SVHN model not found. Falling back to MNIST.")
                self.use_svhn = False
        
        if not self.use_svhn:
            try:
                self.digit_model = self._create_digit_model()
            except Exception as e:
                print(f"Error creating MNIST model: {e}")
                self.digit_model = None
        
        self.categories = ['animals', 'shapes', 'numbers', 'objects']
        self.animal_groups = {
            'bird': ['brambling', 'junco', 'indigo_bunting', 'robin', 'goldfinch', 'house_finch', 
                    'crow', 'raven', 'jay', 'magpie', 'chickadee', 'water_ouzel', 'kite', 'hawk',
                    'eagle', 'vulture', 'peacock', 'parrot', 'macaw', 'cockatoo', 'lorikeet',
                    'hummingbird', 'toucan', 'drake', 'goose', 'swan', 'duck', 'owl', 'ostrich',
                    'pelican', 'king_penguin', 'albatross', 'cock', 'hen', 'quail', 'partridge'],
            'dog': ['chihuahua', 'japanese_spaniel', 'maltese_dog', 'pekinese', 'shih_tzu', 'terrier',
                   'poodle', 'schnauzer', 'labrador', 'retriever', 'spaniel', 'setter', 'sheepdog',
                   'collie', 'border_collie', 'bouvier', 'rottweiler', 'german_shepherd', 'boxer',
                   'bulldog', 'mastiff', 'husky', 'malamute', 'dalmatian', 'affenpinscher', 'pug',
                   'pomeranian', 'chow', 'keeshond', 'dingo', 'dhole', 'african_hunting_dog'],
            
            'cat': ['tabby', 'tiger_cat', 'persian_cat', 'siamese_cat', 'egyptian_cat', 'lion',
                   'tiger', 'cheetah', 'leopard', 'snow_leopard', 'jaguar', 'lynx', 'cougar'],
            
            'large_animal': ['ox', 'bull', 'cow', 'water_buffalo', 'bison', 'horse', 'zebra', 
                            'elephant', 'rhinoceros', 'hippopotamus', 'camel', 'llama', 'deer',
                            'elk', 'moose', 'giraffe', 'bear', 'polar_bear', 'grizzly'],
            
            'small_animal': ['mouse', 'rat', 'hamster', 'gerbil', 'rabbit', 'hare', 'squirrel',
                           'marmot', 'beaver', 'porcupine', 'guinea_pig', 'fox', 'weasel',
                           'mink', 'otter', 'skunk', 'badger', 'ferret', 'mongoose'],
            
            'reptile': ['snake', 'lizard', 'chameleon', 'iguana', 'gecko', 'monitor', 'gila_monster',
                       'alligator', 'crocodile', 'turtle', 'tortoise', 'terrapin'],
            
            'water_animal': ['fish', 'shark', 'ray', 'goldfish', 'carp', 'eel', 'jellyfish',
                           'starfish', 'sea_urchin', 'sea_cucumber', 'dolphin', 'whale', 'seal',
                           'sea_lion', 'walrus', 'octopus', 'squid', 'lobster', 'crab', 'crayfish'],
            
            'insect': ['butterfly', 'moth', 'bee', 'wasp', 'ant', 'fly', 'mosquito', 'dragonfly',
                      'grasshopper', 'cricket', 'beetle', 'ladybug', 'mantis', 'cicada', 'spider'],
            
            'primate': ['monkey', 'marmoset', 'baboon', 'macaque', 'langur', 'colobus', 'proboscis_monkey',
                       'lemur', 'chimpanzee', 'gorilla', 'orangutan', 'gibbon', 'siamang']
        }

    def _create_digit_model(self):
        model = tf.keras.Sequential([
            tf.keras.layers.Flatten(input_shape=(28, 28)),
            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.Dense(10, activation='softmax')
        ])
        
        (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
        x_train = x_train / 255.0
        
        model.compile(optimizer='adam',
                     loss='sparse_categorical_crossentropy',
                     metrics=['accuracy'])
        
        model.fit(x_train[:5000], y_train[:5000], epochs=5, verbose=0)
        return model

    def classify(self, image_path):
        img_color = cv2.imread(image_path)
        
        general_result = self._check_general(image_path)
        if general_result['label'] == "wall clock":
            return {
                'category': 'numbers',
                'detail': "Clock (treated as number)",
                'confidence': general_result['confidence']
            }
        animal_type = self._get_animal_type(general_result['label'])
        shape_result = self._check_shape_improved(img_color)
        number_result = self._check_number(img_color)
        
        if animal_type and general_result['confidence'] > 0.25:
            return {
                'category': 'animals',
                'detail': animal_type,
                'specific': general_result['label'],
                'confidence': general_result['confidence']
            }
        
        if shape_result['confidence'] > 0.6:
            if number_result['confidence'] > 0.95:
                pass
            else:
                return {
                    'category': 'shapes',
                    'detail': shape_result['shape'],
                    'confidence': shape_result['confidence']
                }
        
        if number_result['confidence'] > 0.4:
            return {
                'category': 'numbers',
                'detail': f"Number {number_result['digit']}",
                'confidence': number_result['confidence']
            }
        
        if general_result['confidence'] > 0.4:
            return {
                'category': 'objects',
                'detail': general_result['label'],
                'confidence': general_result['confidence']
            }
        
        candidates = []
        
        if animal_type:
            candidates.append(('animals', animal_type, general_result['confidence'], general_result['label']))
        if shape_result['shape']:
            candidates.append(('shapes', shape_result['shape'], shape_result['confidence'], None))
        if number_result['digit'] is not None:
            candidates.append(('numbers', f"Number {number_result['digit']}", number_result['confidence'], None))
        if not animal_type and general_result['confidence'] > 0.1:
            candidates.append(('objects', general_result['label'], general_result['confidence'], None))
        
        if candidates:
            candidates.sort(key=lambda x: x[2], reverse=True)
            best = candidates[0]
            result = {
                'category': best[0],
                'detail': best[1],
                'confidence': best[2]
            }
            if best[3]:
                result['specific'] = best[3]
            return result
        
        return {'category': 'unknown', 'detail': 'Cannot classify', 'confidence': 0} 

    def _check_number(self, img):
        if self.use_svhn:
            digit, confidence = self.svhn_model.predict_digit(img)
            if digit is not None:
                return {'digit': digit, 'confidence': confidence}
            else:
                return {'digit': None, 'confidence': 0}
        else:
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            resized = cv2.resize(gray, (28, 28))
            
            if np.std(resized) < 30:
                return {'digit': None, 'confidence': 0}
            
            if np.mean(resized) > 127:
                resized = 255 - resized
            
            normalized = resized / 255.0
            prediction = self.digit_model.predict(normalized.reshape(1, 28, 28), verbose=0)[0]
            digit = np.argmax(prediction)
            confidence = prediction[digit]
            
            return {'digit': int(digit), 'confidence': float(confidence)}
    
    def _check_shape(self, img):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        edges = cv2.Canny(blurred, 50, 150)
        result = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        contours = result[-2] 
        
        if not contours:
            return {'shape': None, 'confidence': 0}
        
        contour = max(contours, key=cv2.contourArea)
        epsilon = 0.02 * cv2.arcLength(contour, True)
        approx = cv2.approxPolyDP(contour, epsilon, True)
        vertices = len(approx)
        
        if vertices == 3:
            return {'shape': 'triangle', 'confidence': 0.9}
        elif vertices == 4:
            x, y, w, h = cv2.boundingRect(approx)
            aspect_ratio = float(w) / h
            if 0.9 <= aspect_ratio <= 1.1:
                return {'shape': 'square', 'confidence': 0.9}
            else:
                return {'shape': 'rectangle', 'confidence': 0.85}
        elif vertices >= 8:
            return {'shape': 'circle', 'confidence': 0.85}
        
        return {'shape': None, 'confidence': 0}
    
    def _check_general(self, image_path):
        if not self.general_model_loaded or self.general_model is None:
            return {'label': 'unknown', 'confidence': 0}
        
        try:
            img = Image.open(image_path).convert('RGB')
            img = img.resize((224, 224))
            img_array = np.array(img)
            img_array = preprocess_input(img_array)
            img_array = np.expand_dims(img_array, 0)
            
            predictions = self.general_model.predict(img_array, verbose=0)
            decoded = decode_predictions(predictions, top=1)[0][0]
            
            return {
                'label': decoded[1].replace('_', ' '),
                'confidence': float(decoded[2])
            }
        except Exception as e:
            print(f"Error in general classification: {e}")
            return {'label': 'unknown', 'confidence': 0}
    
    def _get_animal_type(self, specific_label):
        label_lower = specific_label.lower().replace(' ', '_')
        
        for animal_type, animals in self.animal_groups.items():
            if any(animal in label_lower for animal in animals):
                return animal_type
        
        return None

    def _check_shape_improved(self, img):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        if self._looks_like_number(gray):
            return {'shape': None, 'confidence': 0}
        
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        edges = cv2.Canny(blurred, 50, 150)
        result = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        contours = result[-2] 
        
        if not contours:
            return {'shape': None, 'confidence': 0}
        
        contour = max(contours, key=cv2.contourArea)
        
        x, y, w, h = cv2.boundingRect(contour)
        roi = gray[y:y+h, x:x+w]
        filled_pixels = np.sum(roi < 128)
        total_pixels = w * h
        fill_ratio = filled_pixels / total_pixels
        
        if fill_ratio < 0.3:
            return {'shape': None, 'confidence': 0}
        
        epsilon = 0.02 * cv2.arcLength(contour, True)
        approx = cv2.approxPolyDP(contour, epsilon, True)
        vertices = len(approx)
        
        if vertices == 3 and fill_ratio > 0.5:
            return {'shape': 'triangle', 'confidence': 0.9}
        elif vertices == 4 and fill_ratio > 0.6:
            x, y, w, h = cv2.boundingRect(approx)
            aspect_ratio = float(w) / h
            if 0.9 <= aspect_ratio <= 1.1:
                return {'shape': 'square', 'confidence': 0.9}
            else:
                return {'shape': 'rectangle', 'confidence': 0.85}
        elif vertices >= 10 and fill_ratio > 0.7:
            area = cv2.contourArea(contour)
            perimeter = cv2.arcLength(contour, True)
            circularity = 4 * np.pi * area / (perimeter * perimeter)
            if circularity > 0.8:
                return {'shape': 'circle', 'confidence': 0.85}
        
        return {'shape': None, 'confidence': 0}
    
    def _looks_like_number(self, gray_img):
        edges = cv2.Canny(gray_img, 50, 150)
        edge_density = np.sum(edges > 0) / edges.size
        
        if 0.05 < edge_density < 0.3:
            return True
        
        result = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        contours = result[-2] 
        if contours:
            largest = max(contours, key=cv2.contourArea)
            x, y, w, h = cv2.boundingRect(largest)
            aspect_ratio = h / w if w > 0 else 0
            
            if 1.2 < aspect_ratio < 3.0:
                return True
        
        return False
    
    def classify_with_priorities(self, image_path):
        img_color = cv2.imread(image_path)
        
        number_result = self._check_number(img_color)
        if number_result['confidence'] > 0.5:
            return {
                'category': 'numbers',
                'detail': f"Number {number_result['digit']}",
                'confidence': number_result['confidence']
            }
        
        general_result = self._check_general(image_path)
        
        animal_type = self._get_animal_type(general_result['label'])
        if animal_type and general_result['confidence'] > 0.3:
            return {
                'category': 'animals',
                'detail': animal_type,
                'specific': general_result['label'],
                'confidence': general_result['confidence']
            }
        
        if general_result['confidence'] > 0.6:
            return {
                'category': 'objects',
                'detail': general_result['label'],
                'confidence': general_result['confidence']
            }
        
        shape_result = self._check_shape_improved(img_color)
        if shape_result['confidence'] > 0.8:
            return {
                'category': 'shapes',
                'detail': shape_result['shape'],
                'confidence': shape_result['confidence']
            }
        if number_result['confidence'] > 0.3:
            return {
                'category': 'numbers',
                'detail': f"Number {number_result['digit']}",
                'confidence': number_result['confidence']
            }
        
        return {'category': 'unknown', 'detail': 'Cannot classify', 'confidence': 0}