fix v1/detect

main.py

@@ -5,9 +5,9 @@ import aiohttp
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.responses import JSONResponse
 
-# from transformers import AutoImageProcessor, ViTForImageClassification
-from transformers import pipeline
-from transformers.pipelines import PipelineException
+from transformers import AutoImageProcessor, ViTForImageClassification
+# from transformers import pipeline
+# from transformers.pipelines import PipelineException
 from PIL import Image
 from cachetools import Cache
 import torch
@@ -27,10 +27,9 @@ logging.basicConfig(
 cache = Cache(maxsize=1000)
 
 # Load the model using the transformers pipeline
-model = pipeline("image-classification", model="dima806/deepfake_vs_real_image_detection")
-# model = pipeline("image-classification", model="Wvolf/ViT_Deepfake_Detection")
-# image_processor = AutoImageProcessor.from_pretrained("Wvolf/ViT_Deepfake_Detection")
-# model = ViTForImageClassification.from_pretrained("Wvolf/ViT_Deepfake_Detection")
+# model = pipeline("image-classification", model="dima806/deepfake_vs_real_image_detection")
+image_processor = AutoImageProcessor.from_pretrained("dima806/deepfake_vs_real_image_detection")
+model = ViTForImageClassification.from_pretrained("dima806/deepfake_vs_real_image_detection")
     
 # Detect the device used by TensorFlow
 # DEVICE = "GPU" if tf.config.list_physical_devices("GPU") else "CPU"
@@ -84,16 +83,26 @@ async def classify_image(file: UploadFile = File(None)):
 
         image = Image.open(io.BytesIO(image_data))
 
-        # inputs = image_processor(image, return_tensors="pt")
-        inputs = model(image)
-
-        # with torch.no_grad():
-        #     outpus = model(**inputs)
-        #     logits = outpus.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()
+        inputs = image_processor(image, return_tensors="pt")
+        # inputs = model(image)
+
+        with torch.no_grad():
+            outpus = model(**inputs)
+        logits = outpus.logits
+        logging.info("logits %s", logits)
+        probs = F.softmax(logits, dim=-1)
+        logging.info("probs %s", probs)
+        predicted_label_id = probs.argmax(-1).item()
+        logging.info("predicted_label_id %s", predicted_label_id)
+        predicted_label = model.config.id2label[predicted_label_id]
+        logging.info("model.config.id2label %s", model.config.id2label)
+        confidence = probs.max().item()
+# outpus = model(**inputs)
+#             logits = outpus.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()
@@ -101,16 +110,16 @@ async def classify_image(file: UploadFile = File(None)):
     #     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
-        predicted_label = max(inputs, key=lambda x: x["score"])
+        # predicted_label = max(inputs, 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 = round(predicted_label["score"] * 100, 1)
+        # confidence = round(predicted_label["score"] * 100, 1)
 
     # # Prepare the custom response data
-        detection_result = {
+        response_data = {
             "prediction": predicted_label,
             "confidence":confidence,
         }
@@ -130,20 +139,20 @@ async def classify_image(file: UploadFile = File(None)):
         # }
 
         # Populate hash
-        cache[image_hash] = detection_result.copy()
+        cache[image_hash] = response_data.copy()
 
         # Add url to the API response
-        detection_result["file_name"] = file.filename
+        response_data["file_name"] = file.filename
 
-        response_data.append(detection_result)
+        # response_data.append(detection_result)
 
         # Add file_name to the API response
-        response_data["file_name"] = file.filename
+        # response_data["file_name"] = file.filename
 
         return FileImageDetectionResponse(**response_data)
 
-    # except Exception as e:
-    except PipelineException as e:
+    except Exception as e:
+    # except PipelineException as e:
         logging.error("Error processing image: %s", str(e))
         raise HTTPException(
             status_code=500, detail=f"Error processing image: {str(e)}"
@@ -172,29 +181,29 @@ async def classify_images(request: ImageUrlsRequest):
                 continue
 
             image = Image.open(io.BytesIO(image_data))
-            # inputs = image_processor(image, return_tensors="pt")
-            inputs = model(image)
-
-
-            # with torch.no_grad():
-            #     outpus = model(**inputs)
-            # logits = outpus.logits
-            # logging.info("logits %s", logits)
-            # probs = F.softmax(logits, dim=-1)
-            # logging.info("probs %s", probs)
-            # predicted_label_id = probs.argmax(-1).item()
-            # logging.info("predicted_label_id %s", predicted_label_id)
-            # predicted_label = model.config.id2label[predicted_label_id]
-            # logging.info("model.config.id2label %s", model.config.id2label)
-            # confidence = probs.max().item()
+            inputs = image_processor(image, return_tensors="pt")
+            # inputs = model(image)
+
+
+            with torch.no_grad():
+                outpus = model(**inputs)
+            logits = outpus.logits
+            logging.info("logits %s", logits)
+            probs = F.softmax(logits, dim=-1)
+            logging.info("probs %s", probs)
+            predicted_label_id = probs.argmax(-1).item()
+            logging.info("predicted_label_id %s", predicted_label_id)
+            predicted_label = model.config.id2label[predicted_label_id]
+            logging.info("model.config.id2label %s", model.config.id2label)
+            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])
-            logging.info("inputs %s", inputs)
-            predicted_label = max(inputs, key=lambda x: x["score"])
+            # logging.info("inputs %s", inputs)
+            # predicted_label = max(inputs, key=lambda x: x["score"])
     # best_prediction = max(results, key=lambda x: x["score"])
         # logging.info("best_prediction %s", best_prediction)
         # best_prediction2 = results[1]["label"]
@@ -202,7 +211,7 @@ async def classify_images(request: ImageUrlsRequest):
 
         # # Calculate the confidence score, rounded to the nearest tenth and as a percentage
             # confidence_percentage = round(best_prediction["score"] * 100, 1)
-            confidence = round(predicted_label["score"] * 100, 1)
+            # confidence = round(predicted_label["score"] * 100, 1)
 
         # # Prepare the custom response data
             detection_result = {
@@ -232,8 +241,8 @@ async def classify_images(request: ImageUrlsRequest):
 
             response_data.append(detection_result)
 
-        # except Exception as e:
-        except PipelineException as e:
+        except Exception as e:
+        # except PipelineException as e:
             logging.error("Error processing image from %s: %s", image_url, str(e))
             raise HTTPException(
                 status_code=500,

models.py

@@ -23,8 +23,8 @@ class ImageDetectionResponse(BaseModel):
         confidence_percentage (float): Confidence level of the NSFW classification.
     """
 
-    is_nsfw: bool
-    confidence_percentage: float
+    prediction: str
+    confidence: float
 
 
 class FileImageDetectionResponse(ImageDetectionResponse):