remove tf-keras

main.py

@@ -11,7 +11,6 @@ from transformers.pipelines import PipelineException
 from transformers import AutoImageProcessor, ViTForImageClassification
 from PIL import Image
 from cachetools import Cache
-import tensorflow as tf
 import torch
 import torch.nn.functional as F
 from models import (
@@ -34,12 +33,12 @@ cache = Cache(maxsize=1000)
 # model = pipeline("image-classification", model="Wvolf/ViT_Deepfake_Detection")
 
 # Detect the device used by TensorFlow
-DEVICE = "GPU" if tf.config.list_physical_devices("GPU") else "CPU"
-logging.info("TensorFlow version: %s", tf.__version__)
-logging.info("Model is using: %s", DEVICE)
+# DEVICE = "GPU" if tf.config.list_physical_devices("GPU") else "CPU"
+# logging.info("TensorFlow version: %s", tf.__version__)
+# logging.info("Model is using: %s", DEVICE)
 
-if DEVICE == "GPU":
-    logging.info("GPUs available: %d", len(tf.config.list_physical_devices("GPU")))
+# if DEVICE == "GPU":
+#     logging.info("GPUs available: %d", len(tf.config.list_physical_devices("GPU")))
 
 
 async def download_image(image_url: str) -> bytes:
@@ -85,38 +84,58 @@ async def classify_image(file: UploadFile = File(None)):
 
         image = Image.open(io.BytesIO(image_data))
 
-        # Use the model to classify the image
-        # results = model(image)
-
-        image_processor = AutoImageProcessor.from_pretrained("dima806/ai_vs_real_image_detection")
-        model = ViTForImageClassification.from_pretrained("dima806/ai_vs_real_image_detection")
-
+        image_processor = AutoImageProcessor.from_pretrained("Wvolf/ViT_Deepfake_Detection")
+        model = ViTForImageClassification.from_pretrained("Wvolf/ViT_Deepfake_Detection")
         inputs = image_processor(image, return_tensors="pt")
 
         with torch.no_grad():
             logits = model(**inputs).logits
+            probs = F.softmax(logits, dim=-1)
+            predicted_label_id = probs.argmax(-1).item()
+            predicted_label = model.config.id2label[predicted_label_id]
+            confidence = probs.max().item()
+
+    # model predicts one of the 1000 ImageNet classes
+    #     predicted_label = logits.argmax(-1).item()
+    #     logging.info("predicted_label", predicted_label)
+    #     logging.info("model.config.id2label[predicted_label] %s", model.config.id2label[predicted_label])
+    # # print(model.config.id2label[predicted_label])
+    # Find the prediction with the highest confidence using the max() function
+    # best_prediction = max(results, key=lambda x: x["score"])
+    # logging.info("best_prediction %s", best_prediction)
+    # best_prediction2 = results[1]["label"]
+    # logging.info("best_prediction2 %s", best_prediction2)
+
+    # # Calculate the confidence score, rounded to the nearest tenth and as a percentage
+    # confidence_percentage = round(best_prediction["score"] * 100, 1)
+
+    # # Prepare the custom response data
+        detection_result = {
+            "prediction": predicted_label,
+            "confidence_percentage":confidence,
+        }
+        # Use the model to classify the image
+        # results = model(image)
 
-        # model predicts one of the 1000 ImageNet classes
-        predicted_label = logits.argmax(-1).item()
-        logging.info("model.config.id2label[predicted_label] %s", model.config.id2label[predicted_label])
-        # print(model.config.id2label[predicted_label])
         # Find the prediction with the highest confidence using the max() function
         # best_prediction = max(results, key=lambda x: x["score"])
-        # logging.info("best_prediction %s", best_prediction)
-        # best_prediction2 = results[1]["label"]
-        # logging.info("best_prediction2 %s", best_prediction2)
 
-        # # Calculate the confidence score, rounded to the nearest tenth and as a percentage
+        # Calculate the confidence score, rounded to the nearest tenth and as a percentage
         # confidence_percentage = round(best_prediction["score"] * 100, 1)
 
-        # # Prepare the custom response data
-        response_data = {
-            "prediction": model.config.id2label[predicted_label],
-            "confidence_percentage":model.config.id2label[predicted_label],
-        }
+        # Prepare the custom response data
+        # detection_result = {
+        #     "is_nsfw": best_prediction["label"] == "nsfw",
+        #     "confidence_percentage": confidence_percentage,
+        # }
 
         # Populate hash
-        cache[image_hash] = response_data.copy()
+        cache[image_hash] = detection_result.copy()
+
+        # Add url to the API response
+        detection_result["file_name"] = file.filename
+
+        response_data.append(detection_result)
 
         # Add file_name to the API response
         response_data["file_name"] = file.filename
@@ -214,10 +233,6 @@ async def classify_images(request: ImageUrlsRequest):
 
     return JSONResponse(status_code=200, content=response_data)
 
-@app.get("/hello")
-async def hello_world():
-   return {"message": "hello_world"}
-
 if __name__ == "__main__":
     import uvicorn
 

requirements.txt

@@ -5,6 +5,5 @@ aiohttp==3.9.5
 pillow==10.3.0
 python-multipart==0.0.9
 torch
-tf-keras==2.16.0
 cachetools===5.3.3
 pydantic===2.7.2