Update app.py

app.py

@@ -1,67 +1,57 @@
-import os
-import requests
-from datasets import load_dataset
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
-
-# Step 1: Set up environment and check paths
-MODEL_PATH = "yolac/BacterialMorphologyClassification"
-DATASET_PATH = "yolac/BacterialMorphologyClassification"
-
-# Step 2: Load the model
-try:
-    model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH, num_labels=3)
-    print("Model loaded successfully.")
-except Exception as e:
-    print(f"Failed to load the model. Error: {e}")
-
-# Step 3: Load the tokenizer (if needed)
-try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-    print("Tokenizer loaded successfully.")
-except Exception as e:
-    print(f"Failed to load the tokenizer. Error: {e}")
-
-# Step 4: Load the dataset
-try:
-    dataset = load_dataset(DATASET_PATH, split="train")
-    print("Dataset loaded successfully.")
-except Exception as e:
-    print(f"Failed to load the dataset. Error: {e}")
-
-# Step 5: Preprocess and prepare data for model input (example code)
-def preprocess_data(example):
-    # Add any necessary preprocessing steps here, e.g., tokenization
-    return tokenizer(example['text'], padding="max_length", truncation=True)
-
-# Apply preprocessing
-dataset = dataset.map(preprocess_data, batched=True)
-
-# Step 6: Set up training arguments (use the `Trainer` class if needed)
-from transformers import Trainer, TrainingArguments
-
-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=3,
-    weight_decay=0.01,
-    logging_dir="./logs",
-)
-
-# Initialize the Trainer
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=dataset,
-    tokenizer=tokenizer,
-)
-
-# Step 7: Train the model
-try:
-    trainer.train()
-    print("Training completed successfully.")
-except Exception as e:
-    print(f"Failed to train the model. Error: {e}")
+from fastapi import FastAPI, UploadFile, File
+from pydantic import BaseModel
+from io import BytesIO
+from PIL import Image
+from torchvision import transforms
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
+
+# Initialize the FastAPI app
+app = FastAPI()
+
+# Load the model and feature extractor
+model = AutoModelForImageClassification.from_pretrained("yolac/BacterialMorphologyClassification")
+feature_extractor = AutoFeatureExtractor.from_pretrained("yolac/BacterialMorphologyClassification")
+
+# Define the image preprocessing transform
+preprocess = transforms.Compose([
+    transforms.Resize((224, 224)),  # Adjust size as needed for your model
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+# Helper function to preprocess the image
+def preprocess_image(image: Image.Image):
+    image = image.convert("RGB")  # Ensure the image is in RGB format
+    image = preprocess(image)
+    image = image.unsqueeze(0)  # Add a batch dimension
+    return image
+
+# Define the prediction endpoint
+@app.post("/predict/")
+async def predict(file: UploadFile = File(...)):
+    try:
+        # Read the uploaded image
+        image_data = await file.read()
+        image = Image.open(BytesIO(image_data))
+
+        # Preprocess the image
+        image_tensor = preprocess_image(image)
+
+        # Perform inference
+        model.eval()
+        with torch.no_grad():
+            outputs = model(image_tensor)
+            logits = outputs.logits  # Get the model's raw output
+            predicted_class_idx = torch.argmax(logits, dim=1).item()
+
+        # Map the predicted class index to the class labels
+        class_labels = ["Cocci", "Bacilli", "Spirilla"]  # Replace with your actual class labels
+        predicted_class = class_labels[predicted_class_idx]
+
+        return {"predicted_class": predicted_class}
+
+    except Exception as e:
+        return {"error": str(e)}
+
+# Run the app using the command: uvicorn app:app --reload