Update app.py

app.py

@@ -1,23 +1,25 @@
 import gradio as gr
-import requests
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
+import torch
 from PIL import Image
 
-def load_model(model_name):
-    # Use the Hugging Face API to load the model
-    api_url = f"https://huggingface.co/{model_name}"
-    response = requests.get(api_url)
-    if response.status_code == 200:
-        # Assume the model is an image processing model
-        return lambda image: image  # Replace with actual model processing code
-    else:
-        raise ValueError("Model not found")
-
-# Load the model once during initialization
-model = load_model("amjadfqs/finalProject")
+# Load the model and feature extractor once during initialization
+model_name = "amjadfqs/finalProject"
+model = AutoModelForImageClassification.from_pretrained(model_name)
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
 
 def predict(image):
-    # Use the model to make a prediction
-    return model(image)
+    # Preprocess the image
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    # Make prediction
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.logits
+    # Get the predicted class
+    predicted_class = logits.argmax(-1).item()
+    # You may need to adjust the following line based on your class labels
+    class_names = ["class1", "class2", "class3", "class4"]
+    return predicted_class
 
 # Set up the Gradio interface
 image_cp = gr.Image(type="pil", label='Brain')