Update app.py

app.py

@@ -15,6 +15,7 @@ from PIL import Image
 from ultralytics import YOLO
 from PIL import Image
 from facenet_pytorch import MTCNN
+from ultralytics import YOLO
 
 
 app = Flask(__name__)
@@ -25,6 +26,8 @@ binary_labels = ['Real','Spoof']
 multi_voice_labels = ['Real','Text to Speech','Voice Conversion','Text to Speech + Voice Conversion']
 multi_face_labels = ['Genuine Face','Printed Photo','Paper Cut','Replayed Face','3D Mask']
 
+finger_print_detector = YOLO('fingerprint_best.pt')
+
 
 def process_audio(encoded_audio):
     decoded_audio = base64.b64decode(encoded_audio)
@@ -83,6 +86,7 @@ def process_image(base64_img,extend=0):
 def process_fingerprint_image(base64_img):
     image_data = base64.b64decode(base64_img)
     img = Image.open(io.BytesIO(image_data)).convert('RGB')
+    
     if isinstance(img, torch.Tensor):
         img = transforms.ToPILImage()(img)
     elif isinstance(img, np.ndarray):
@@ -94,7 +98,9 @@ def process_fingerprint_image(base64_img):
         ])
         
     img = transformer(img)
-    return img.unsqueeze(0)
+    results = finger_print_detector(img)
+    is_detected = any(np.array(results[0].boxes.cls.cpu())==0)
+    return img.unsqueeze(0),is_detected
 
 
 
@@ -462,22 +468,25 @@ def post_api_fingerprint():
         encoded_image = data['base64']
         
         # Process the image 
-        processsed_image = process_fingerprint_image(encoded_image)
-        
-        with torch.no_grad():
-            output, _ = vit_binary_model(processsed_image)
-            prob = torch.nn.functional.softmax(output[0], dim=0)
-            pred = torch.argmax(prob).item()
-            category = binary_labels[pred]
-            probs_dict = {binary_labels[i]: prob[i].item()*100 for i in range(len(binary_labels))}
+        processsed_image, is_detected = process_fingerprint_image(encoded_image)
 
-        response = {
-            'message': 'Data received!',
-            'class': category,
-            'probs': probs_dict,
-            'status':200
-        }
-        return jsonify(response), 201
+        if is_detected:
+            with torch.no_grad():
+                output, _ = vit_binary_model(processsed_image)
+                prob = torch.nn.functional.softmax(output[0], dim=0)
+                pred = torch.argmax(prob).item()
+                category = binary_labels[pred]
+                probs_dict = {binary_labels[i]: prob[i].item()*100 for i in range(len(binary_labels))}
+    
+            response = {
+                'message': 'Data received!',
+                'class': category,
+                'probs': probs_dict,
+                'status':200
+            }
+            return jsonify(response), 201
+        else:
+            return jsonify({'error': f'No Fingerprint Detected','status':400}), 400
     
     except KeyError as e:
         return jsonify({'error': f'Missing key: {str(e)}','status':400}), 400