added api-keys

main.py

@@ -1,80 +1,86 @@
-import torch
-import torch.nn.functional as F
-from facenet_pytorch import MTCNN, InceptionResnetV1
-import numpy as np
-from PIL import Image
-import cv2
-from pytorch_grad_cam import GradCAM
-from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
-from pytorch_grad_cam.utils.image import show_cam_on_image
-from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import JSONResponse
-from io import BytesIO
-
-app = FastAPI()
-
-DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
-
-mtcnn = MTCNN(select_largest=False, post_process=False, device=DEVICE).eval()
-
-model = InceptionResnetV1(classify=True, num_classes=1, device=DEVICE)
-model_path = "Model/resnetinceptionv1_epoch_32.pth"
-checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
-model.load_state_dict(checkpoint['model_state_dict'])
-model.to(DEVICE)
-model.eval()
-
-def predict(input_image: Image.Image):
-    """Predict the label of the input_image"""
-    if input_image.mode == 'RGBA':
-        input_image = input_image.convert('RGB')
-    face = mtcnn(input_image)
-    if face is None:
-        raise Exception('No face detected')
-    face = face.unsqueeze(0)  # add the batch dimension
-    face = F.interpolate(face, size=(256, 256), mode='bilinear', align_corners=False)
-
-    prev_face = face.squeeze(0).permute(1, 2, 0).cpu().detach().int().numpy()
-    prev_face = prev_face.astype('uint8')
-
-    face = face.to(DEVICE)
-    face = face.to(torch.float32)
-    face = face / 255.0
-    face_image_to_plot = face.squeeze(0).permute(1, 2, 0).cpu().detach().int().numpy()
-
-    target_layers = [model.block8.branch1[-1]]
-    cam = GradCAM(model=model, target_layers=target_layers)
-    targets = [ClassifierOutputTarget(0)]
-
-    grayscale_cam = cam(input_tensor=face, targets=targets, eigen_smooth=True)
-    grayscale_cam = grayscale_cam[0, :]
-    visualization = show_cam_on_image(face_image_to_plot, grayscale_cam, use_rgb=True)
-    face_with_mask = cv2.addWeighted(prev_face, 1, visualization, 0.5, 0)
-
-    with torch.no_grad():
-        output = torch.sigmoid(model(face).squeeze(0))
-        prediction = "real" if output.item() < 0.5 else "fake"
-
-        real_prediction = 1 - output.item()
-        fake_prediction = output.item()
-
-        confidences = {
-            'real': real_prediction,
-            'fake': fake_prediction
-        }
-    return confidences, prediction, face_with_mask
-
-@app.post("/predict")
-async def predict_api(file: UploadFile = File(...)):
-    image = Image.open(BytesIO(await file.read()))
-    try:
-        confidences, prediction, face_with_mask = predict(image)
-        _, buffer = cv2.imencode('.jpg', face_with_mask)
-        face_with_mask_encoded = buffer.tobytes()
-        return JSONResponse(content={
-            "confidences": confidences,
-            "prediction": prediction,
-            "face_with_mask": face_with_mask_encoded.hex()
-        })
-    except Exception as e:
-        return JSONResponse(content={"error": str(e)}, status_code=400)
+import torch
+import torch.nn.functional as F
+from facenet_pytorch import MTCNN, InceptionResnetV1
+import numpy as np
+from PIL import Image
+import cv2
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+from pytorch_grad_cam.utils.image import show_cam_on_image
+from fastapi import FastAPI, File, UploadFile, Depends, HTTPException, Header
+from fastapi.responses import JSONResponse
+from io import BytesIO
+
+app = FastAPI()
+
+DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+
+mtcnn = MTCNN(select_largest=False, post_process=False, device=DEVICE).eval()
+
+model = InceptionResnetV1(classify=True, num_classes=1, device=DEVICE)
+model_path = "Model/resnetinceptionv1_epoch_32.pth"
+checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
+model.load_state_dict(checkpoint['model_state_dict'])
+model.to(DEVICE)
+model.eval()
+
+API_KEY = "c50dd5ady0uRL0rdnSaVyrArYaN161edb06af8"
+
+def get_api_key(api_key: str = Header(...)):
+    if api_key != API_KEY:
+        raise HTTPException(status_code=403, detail="Could not validate credentials")
+
+def predict(input_image: Image.Image):
+    """Predict the label of the input_image"""
+    if input_image.mode == 'RGBA':
+        input_image = input_image.convert('RGB')
+    face = mtcnn(input_image)
+    if face is None:
+        raise Exception('No face detected')
+    face = face.unsqueeze(0)  # add the batch dimension
+    face = F.interpolate(face, size=(256, 256), mode='bilinear', align_corners=False)
+
+    prev_face = face.squeeze(0).permute(1, 2, 0).cpu().detach().int().numpy()
+    prev_face = prev_face.astype('uint8')
+
+    face = face.to(DEVICE)
+    face = face.to(torch.float32)
+    face = face / 255.0
+    face_image_to_plot = face.squeeze(0).permute(1, 2, 0).cpu().detach().int().numpy()
+
+    target_layers = [model.block8.branch1[-1]]
+    cam = GradCAM(model=model, target_layers=target_layers)
+    targets = [ClassifierOutputTarget(0)]
+
+    grayscale_cam = cam(input_tensor=face, targets=targets, eigen_smooth=True)
+    grayscale_cam = grayscale_cam[0, :]
+    visualization = show_cam_on_image(face_image_to_plot, grayscale_cam, use_rgb=True)
+    face_with_mask = cv2.addWeighted(prev_face, 1, visualization, 0.5, 0)
+
+    with torch.no_grad():
+        output = torch.sigmoid(model(face).squeeze(0))
+        prediction = "real" if output.item() < 0.5 else "fake"
+
+        real_prediction = 1 - output.item()
+        fake_prediction = output.item()
+
+        confidences = {
+            'real': real_prediction,
+            'fake': fake_prediction
+        }
+    return confidences, prediction, face_with_mask
+
+@app.post("/predict")
+async def predict_api(file: UploadFile = File(...), api_key: str = Depends(get_api_key)):
+    image = Image.open(BytesIO(await file.read()))
+    try:
+        confidences, prediction, face_with_mask = predict(image)
+        _, buffer = cv2.imencode('.jpg', face_with_mask)
+        face_with_mask_encoded = buffer.tobytes()
+        return JSONResponse(content={
+            "confidences": confidences,
+            "prediction": prediction,
+            "face_with_mask": face_with_mask_encoded.hex()
+        })
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)}, status_code=400)