Add application file

app.py

@@ -2,16 +2,11 @@ import os
 import warnings
 from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 import torch
-warnings.filterwarnings("ignore")
-
-import json
 from flask_cors import CORS
-from flask import Flask, request, Response
-
+from flask import Flask, request, json, Response
 import numpy as np
 from PIL import Image
 import requests
-
 from io import BytesIO
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -19,36 +14,49 @@ os.environ["CUDA_VISIBLE_DEVICES"] = ""
 app = Flask(__name__)
 cors = CORS(app)
 
-global MODEL
-global CLASSES
-
+# Define the model and feature extractor globally
+model = AutoModelForImageClassification.from_pretrained('carbon225/vit-base-patch16-224-hentai')
+feature_extractor = AutoFeatureExtractor.from_pretrained('carbon225/vit-base-patch16-224-hentai')
 
 @app.route("/", methods=["GET"])
 def default():
     return json.dumps({"Hello I am Chitti": "Speed 1 Terra Hertz, Memory 1 Zeta Byte"})
 
-
 @app.route("/predict", methods=["GET"])
 def predict():
-    feature_extractor = AutoFeatureExtractor.from_pretrained('carbon225/vit-base-patch16-224-hentai')
-    model = AutoModelForImageClassification.from_pretrained('carbon225/vit-base-patch16-224-hentai')
-    src = request.args.get("src")
-    print(f"{src=}")
-    response = requests.get(src)
-    print(f"{response=}")
     try:
+        src = request.args.get("src")
+        print(f"{src=}")
+
+        # Download image from the provided URL
+        response = requests.get(src)
+        response.raise_for_status()  # Check for HTTP errors
+
+        # Open and preprocess the image
         image = Image.open(BytesIO(response.content))
         image = image.resize((128, 128))
-        encoding = feature_extractor(image.convert("RGB"), return_tensors="pt")
+
+        # Extract features using the pre-trained feature extractor
+        encoding = feature_extractor(images=image.convert("RGB"), return_tensors="pt")
+
+        # Make a prediction using the pre-trained model
         with torch.no_grad():
             outputs = model(**encoding)
             logits = outputs.logits
 
+        # Get the predicted class index and label
         predicted_class_idx = logits.argmax(-1).item()
-        print(model.config.id2label[predicted_class_idx])
-        # Return the Predictions
-        return json.dumps({"class": model.config.id2label[predicted_class_idx]})
-    except Exception as e:
-        return json.dumps({"Uh oh": f"{str(e)}"})
+        predicted_class_label = model.config.id2label[predicted_class_idx]
 
+        print(predicted_class_label)
+
+        # Return the predictions
+        return json.dumps({"class": predicted_class_label})
+
+    except requests.exceptions.RequestException as e:
+        return json.dumps({"error": f"Request error: {str(e)}"})
+    except Exception as e:
+        return json.dumps({"error": f"An unexpected error occurred: {str(e)}"})
 
+if __name__ == "__main__":
+    app.run(debug=True)