Update app.py

app.py

@@ -1,143 +1,46 @@
-import os
+import gradio as gr
+import tensorflow as tf
+from PIL import Image
 import numpy as np
-from tensorflow.keras.preprocessing.image import ImageDataGenerator
-from tensorflow.keras.applications import ResNet50
-from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, Dropout, BatchNormalization
-from tensorflow.keras.models import Model
-from tensorflow.keras.optimizers import Adam
-from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau
-import matplotlib.pyplot as plt
-from collections import Counter
-
-# Pfad zum Verzeichnis der Bilddaten
-base_dir = 'pokemon-images-first-generation17000-files/pokemon'  # Pfad anpassen
-
-# Überprüfen, ob das Verzeichnis existiert
-if not os.path.exists(base_dir):
-    raise Exception("Der angegebene Basispfad existiert nicht. Bitte überprüfen Sie den Pfad.")
-
-# Definiere die Pokémon-Klassen, die klassifiziert werden sollen
-classes = ['Dragonite', 'Bulbasaur', 'Golbat']
-
-# Data Augmentation und Daten-Generatoren
-train_datagen = ImageDataGenerator(
-    rescale=1./255,
-    rotation_range=45,
-    width_shift_range=0.25,
-    height_shift_range=0.25,
-    shear_range=0.25,
-    zoom_range=0.3,
-    horizontal_flip=True,
-    fill_mode='nearest',
-    validation_split=0.2)
-
-validation_datagen = ImageDataGenerator(
-    rescale=1./255,
-    validation_split=0.2)
-
-train_generator = train_datagen.flow_from_directory(
-    base_dir,
-    target_size=(150, 150),
-    batch_size=32,
-    classes=classes,
-    class_mode='categorical',
-    subset='training')
-
-validation_generator = validation_datagen.flow_from_directory(
-    base_dir,
-    target_size=(150, 150),
-    batch_size=32,
-    classes=classes,
-    class_mode='categorical',
-    subset='validation')
-
-# Basis-Modell ResNet50 laden und anpassen
-base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(150, 150, 3))
-base_model.trainable = False  # Initial alle Schichten einfrieren
-
-x = base_model.output
-x = GlobalAveragePooling2D()(x)
-x = Dense(1024, activation='relu')(x)
-x = Dropout(0.3)(x)
-x = BatchNormalization()(x)
-predictions = Dense(len(classes), activation='softmax')(x)
-model = Model(inputs=base_model.input, outputs=predictions)
-
-# Callbacks einrichten
-callbacks = [
-    ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5, min_lr=0.00001, verbose=1),
-    EarlyStopping(monitor='val_loss', patience=10, verbose=1),
-    ModelCheckpoint('best_model.keras', monitor='val_loss', save_best_only=True, verbose=1)
-]
-
-# Modell kompilieren mit Adam Optimizer
-model.compile(optimizer=Adam(learning_rate=0.001), loss='categorical_crossentropy', metrics=['accuracy'])
-
-# Diagnostic output
-print(f"Total training samples: {train_generator.samples}")
-print(f"Total validation samples: {validation_generator.samples}")
-print(f"Training batch size: {train_generator.batch_size}")
-print(f"Validation batch size: {validation_generator.batch_size}")
-
-# Calculate steps per epoch and validation steps
-training_steps_per_epoch = int(np.ceil(train_generator.samples / train_generator.batch_size))
-validation_steps_per_epoch = int(np.ceil(validation_generator.samples / validation_generator.batch_size))
-
-print(f"Calculated training steps per epoch: {training_steps_per_epoch}")
-print(f"Calculated validation steps per epoch: {validation_steps_per_epoch}")
-
-# Model training with updated steps per epoch
-history = model.fit(
-    train_generator,
-    steps_per_epoch=training_steps_per_epoch,
-    validation_data=validation_generator,
-    validation_steps=validation_steps_per_epoch,
-    epochs=15,
-    callbacks=callbacks
-)
-
-# Modell speichern
-model.save('pokemon_classifier_model.h5')
-
-# Fine-Tuning: Einige Schichten des Basis-Modells freigeben
-base_model.trainable = True
-for layer in base_model.layers[:165]:
-    layer.trainable = False
-
-# Modell neu kompilieren mit Adam für Fine-Tuning
-model.compile(optimizer=Adam(learning_rate=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])
-
-# Fine-Tuning trainieren mit Callbacks
-fine_tune_history = model.fit(
-    train_generator,
-    steps_per_epoch=train_generator.samples // train_generator.batch_size,
-    validation_data=validation_generator,
-    validation_steps=validation_generator.samples // validation_generator.batch_size,
-    epochs=10,
-    callbacks=callbacks)
-
-# Bilder und deren Labels aus dem Validierungsdatengenerator holen
-class_indices = train_generator.class_indices
-class_names = list(class_indices.keys())
-
-for images, labels in validation_generator:
-    # Nur die ersten 5 Bilder verwenden
-    for i in range(5):
-        img = images[i]
-        label = labels[i]
-        plt.imshow(img)
-        
-        # Vorhersage erstellen
-        img_array = np.expand_dims(img, axis=0)  # Modell erwartet einen Batch
-        predictions = model.predict(img_array)
-        predicted_class_index = np.argmax(predictions, axis=1)[0]
-        predicted_class_name = class_names[predicted_class_index]
-        
-        # Tatsächliche Klasse herausfinden
-        true_class_index = np.argmax(label)
-        true_class_name = class_names[true_class_index]
-        
-        plt.title(f'Vorhergesagt: {predicted_class_name}, Wahr: {true_class_name}')
-        plt.show()
-       
-    break  # Nur die erste Charge von Bildern verwenden
+ 
+# Load the pre-trained Pokémon model
+model_path = "pokemon_classifier_model.keras"
+model = tf.keras.models.load_model(model_path)
+ 
+# Define the Pokémon classes
+classes = ['Doduo', 'Geodude', 'Zubat']  # Adjust these as per your model's classes
+ 
+# Define the image classification function
+def classify_image(image):
+    try:
+        # Ensure the image is in RGB and normalize it
+        if image.ndim == 2:  # Check if the image is grayscale
+            image = np.stack((image,)*3, axis=-1)  # Convert grayscale to RGB by repeating the gray channel
+        elif image.shape[2] == 4:  # Check if the image has an alpha channel
+            image = image[:, :, :3]  # Drop the alpha channel
+        image = Image.fromarray(image.astype('uint8'), 'RGB')  # Convert to PIL Image to resize
+        image = image.resize((150, 150))  # Resize to match the model's input size
+ 
+        image_array = np.array(image) / 255.0  # Convert to array and normalize
+        image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
+ 
+        # Predict using the model
+        prediction = model.predict(image_array)
+        predicted_class = classes[np.argmax(prediction)]
+        confidence = np.max(prediction)
+ 
+        return f"Predicted Pokémon: {predicted_class}, Confidence: {np.round(confidence * 100, 2)}%"
+    except Exception as e:
+        return str(e)  # Return the error message if something goes wrong
+ 
+# Create Gradio interface
+input_image = gr.Image()  # Using Gradio's Image component correctly
+output_label = gr.Label()
+ 
+interface = gr.Interface(fn=classify_image,
+                         inputs=input_image,
+                         outputs=output_label,
+                         examples=["pokemon-images-first-generation17000-files/pokemon/Bulbasaur/00000000.png", "pokemon-images-first-generation17000-files/pokemon/Dragonite/00000000.jpg", "pokemon-images-first-generation17000-files/pokemon/Golbat/00000000.png"],
+                         description="Upload an image of a Pokémon to classify!")
+ 
+interface.launch()