Charlesremador/Judo

import tensorflow as tf from tensorflow.keras.preprocessing.image import ImageDataGenerator from tensorflow.keras.models import Sequential from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint import numpy as np

class ModeloJudo: def init(self, altura_imagem=224, largura_imagem=224, num_classes=50): """ Inicializa o modelo para classificação de golpes de judô. """ self.altura_imagem = altura_imagem self.largura_imagem = largura_imagem self.num_classes = num_classes self.model = None

def criar_geradores_dados(self, diretorio_treino, diretorio_teste, batch_size=32):
    """
    Cria geradores de dados para augmentação e pré-processamento das imagens.
    """
    treino_datagen = ImageDataGenerator(
        rescale=1./255,
        rotation_range=20,
        width_shift_range=0.2,
        height_shift_range=0.2,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True,
        fill_mode='nearest'
    )
    teste_datagen = ImageDataGenerator(rescale=1./255)

    treino_generator = treino_datagen.flow_from_directory(
        diretorio_treino,
        target_size=(self.altura_imagem, self.largura_imagem),
        batch_size=batch_size,
        class_mode='categorical'
    )
    teste_generator = teste_datagen.flow_from_directory(
        diretorio_teste,
        target_size=(self.altura_imagem, self.largura_imagem),
        batch_size=batch_size,
        class_mode='categorical'
    )

    return treino_generator, teste_generator

def criar_modelo(self):
    """
    Cria a arquitetura do modelo CNN para classificação de golpes.
    """
    modelo = Sequential([
        Conv2D(32, (3, 3), activation='relu', input_shape=(self.altura_imagem, self.largura_imagem, 3)),
        MaxPooling2D(2, 2),
        Conv2D(64, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),
        Conv2D(128, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),
        Conv2D(128, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),
        Flatten(),
        Dense(512, activation='relu'),
        Dropout(0.5),
        Dense(self.num_classes, activation='softmax')
    ])

    self.model = modelo
    return modelo

def treinar_modelo(self, treino_generator, teste_generator, epocas=50):
    """
    Treina o modelo usando os geradores de dados fornecidos.
    """
    if self.model is None:
        self.criar_modelo()

    self.model.compile(
        optimizer='adam',
        loss='categorical_crossentropy',
        metrics=['accuracy']
    )

    callbacks = [
        EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True),
        ModelCheckpoint('melhor_modelo_judo.h5', monitor='val_accuracy', save_best_only=True)
    ]

    history = self.model.fit(
        treino_generator,
        epochs=epocas,
        validation_data=teste_generator,
        callbacks=callbacks
    )

    return history

def salvar_modelo(self, caminho='modelo_judo.h5'):
    """
    Salva o modelo treinado.
    """
    if self.model is not None:
        self.model.save(caminho)
    else:
        raise Exception("Modelo ainda não foi criado!")

def carregar_modelo(self, caminho='modelo_judo.h5'):
    """
    Carrega um modelo previamente salvo.
    """
    self.model = tf.keras.models.load_model(caminho)

def prever_golpe(self, caminho_imagem):
    """
    Faz a previsão do golpe para uma única imagem.
    """
    imagem = tf.keras.preprocessing.image.load_img(
        caminho_imagem,
        target_size=(self.altura_imagem, self.largura_imagem)
    )
    imagem = tf.keras.preprocessing.image.img_to_array(imagem)
    imagem = np.expand_dims(imagem, axis=0)
    imagem = imagem / 255.0

    predicao = self.model.predict(imagem)
    classe_prevista = np.argmax(predicao[0])
    probabilidade = predicao[0][classe_prevista]

    return classe_prevista, probabilidade

if name == "main": # Parâmetros NUM_CLASSES = 50 # Número de classes (golpes) DIRETORIO_TREINO = 'dados/treino' DIRETORIO_TESTE = 'dados/teste' CAMINHO_MODELO = 'modelo_judo_treinado.h5' CAMINHO_IMAGEM = 'dados/imagem_teste/judo_exemplo.jpg' # Caminho automático da imagem

# Lista de 50 golpes clássicos do judô
CLASSES = [
    'O-Soto-Gari', 'Ippon-Seoi-Nage', 'O-Goshi', 'Harai-Goshi', 'Uchi-Mata', 'Ko-Uchi-Gari',
    'Tai-Otoshi', 'De-Ashi-Barai', 'Sasae-Tsuri-Komi-Ashi', 'Koshi-Guruma',
    'Hane-Goshi', 'Okuri-Ashi-Barai', 'Tani-Otoshi', 'Hiza-Guruma', 'O-Uchi-Gari',
    'Ushiro-Goshi', 'Yoko-Otoshi', 'Kata-Guruma', 'Ashi-Guruma', 'Soto-Makikomi',
    'Uki-Goshi', 'Morote-Seoi-Nage', 'Tomoe-Nage', 'Yoko-Tomoe-Nage', 'Kubi-Nage',
    'Sumi-Gaeshi', 'Tsurikomi-Goshi', 'Ura-Nage', 'Ushiro-Keshi', 'Utsuri-Goshi',
    'O-Soto-Guruma', 'Harai-Tsurikomi-Ashi', 'Hane-Makikomi', 'Uki-Otoshi', 'Kuchiki-Daoshi',
    'Obi-Tori-Gaeshi', 'Te-Guruma', 'Koshi-Jime', 'Hadaka-Jime', 'Okuri-Eri-Jime',
    'Sode-Tsurikomi-Goshi', 'O-Tsuri-Goshi', 'Seoi-Otoshi', 'Sankaku-Jime', 'Do-Jime',
    'Gyaku-Juji-Jime', 'Katagatame', 'Kesa-Gatame', 'Kami-Shiho-Gatame', 'Tate-Shiho-Gatame'
]

# Criar instância do modelo
modelo_judo = ModeloJudo(num_classes=NUM_CLASSES)

# Criar geradores de dados
treino_gen, teste_gen = modelo_judo.criar_geradores_dados(
    DIRETORIO_TREINO,
    DIRETORIO_TESTE
)

# Treinar o modelo
modelo_judo.criar_modelo()
historico = modelo_judo.treinar_modelo(treino_gen, teste_gen, epocas=50)

# Salvar o modelo treinado
modelo_judo.salvar_modelo(CAMINHO_MODELO)

# Previsão automática com o caminho da imagem
modelo_judo.carregar_modelo(CAMINHO_MODELO)
classe, probabilidade = modelo_judo.prever_golpe(CAMINHO_IMAGEM)
print(f"Classe prevista: {CLASSES[classe]}, Probabilidade: {probabilidade:.2f}")