Update app.py

app.py

@@ -1,10 +1,88 @@
-from flask import Flask, jsonify
+import os
+import numpy as np
+import cv2
+from sklearn.model_selection import train_test_split
+from transformers import ViTForImageClassification, ViTFeatureExtractor
+from transformers import Trainer, TrainingArguments
+from datasets import load_dataset, Dataset
 
-app = Flask(__name__)
+# Function to load images from specified folders
+def load_images_from_folders(folders, label):
+    images = []
+    labels = []
+    for folder in folders:
+        for filename in os.listdir(folder):
+            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):  # Check for valid image extensions
+                img = cv2.imread(os.path.join(folder, filename), cv2.IMREAD_GRAYSCALE)  # Read image as grayscale
+                if img is not None:
+                    img = cv2.resize(img, (224, 224))  # Resize to 224x224 pixels
+                    img = img.astype(np.float32)  # Ensure the image is in float32 format
+                    img /= 255.0  # Normalize to [0, 1]
+                    images.append(img)
+                    labels.append(label)
+                else:
+                    print(f"Failed to load image: {filename}")
+    return images, labels
 
-@app.route('/api/data')
-def get_data():
-    return jsonify({"message": "Hello from Flask!"})
+# Load normal and pneumonia images
+normal_folders = [
+    os.path.join('chest-xray-pneumonia', 'chest_xray', 'test', 'NORMAL'),
+    os.path.join('chest-xray-pneumonia', 'chest_xray', 'train', 'NORMAL'),
+    os.path.join('chest-xray-pneumonia', 'chest_xray', 'val', 'NORMAL'),
+]
+pneumonia_folders = [
+    os.path.join('chest-xray-pneumonia', 'chest_xray', 'test', 'PNEUMONIA'),
+    os.path.join('chest-xray-pneumonia', 'chest_xray', 'train', 'PNEUMONIA'),
+    os.path.join('chest-xray-pneumonia', 'chest_xray', 'val', 'PNEUMONIA'),
+]
 
-if __name__ == "__main__":
-    app.run(port=3000)  # Run on port 5000
+normal_images, normal_labels = load_images_from_folders(normal_folders, 0)
+pneumonia_images, pneumonia_labels = load_images_from_folders(pneumonia_folders, 1)
+
+# Combine images and labels
+images = normal_images + pneumonia_images
+labels = normal_labels + pneumonia_labels
+
+# Split the dataset into training and testing sets
+X_train, X_test, y_train, y_test = train_test_split(images, labels, test_size=0.2, random_state=42)
+
+# Convert the dataset to a Hugging Face Dataset
+train_dataset = Dataset.from_dict({"image": X_train, "label": y_train})
+test_dataset = Dataset.from_dict({"image": X_test, "label": y_test})
+
+# Load feature extractor and model
+feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
+model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k', num_labels=2)
+
+# Preprocess the dataset
+def preprocess_function(examples):
+    return feature_extractor(images=examples['image'], return_tensors="pt")
+
+train_dataset = train_dataset.map(preprocess_function, batched=True)
+test_dataset = test_dataset.map(preprocess_function, batched=True)
+
+# Training arguments
+training_args = TrainingArguments(
+    output_dir='./results',
+    evaluation_strategy='epoch',
+    learning_rate=2e-5,
+    per_device_train_batch_size=8,
+    per_device_eval_batch_size=8,
+    num_train_epochs=10,
+    weight_decay=0.01,
+)
+
+# Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=test_dataset,
+)
+
+# Train the model
+trainer.train()
+
+# Save the model
+model.save_pretrained('./pneumonia_model_final')
+print("Model saved as './pneumonia_model_final'")