Update app.py

app.py

@@ -25,7 +25,10 @@ os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 # Device configuration
 DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
 
-mtcnn = MTCNN(select_largest=False, post_process=False, device=DEVICE).to(DEVICE).eval()
+# Configure MTCNN with adjusted thresholds
+mtcnn = MTCNN(select_largest=False, post_process=False, device=DEVICE,
+              thresholds=[0.7, 0.8, 0.8],  # Adjust these thresholds for P-Net, R-Net, O-Net
+              margin=20, min_face_size=50).to(DEVICE).eval()
 
 model = InceptionResnetV1(pretrained="vggface2", classify=True, num_classes=1, device=DEVICE)
 checkpoint = torch.load("resnetinceptionv1_epoch_32.pth", map_location=torch.device('cpu'))
@@ -41,18 +44,39 @@ targets = [ClassifierOutputTarget(0)]
 def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
+def filter_low_quality_detections(detection, min_size=(50, 50)):
+    if detection is None or detection[0] is None:
+        return None
+    for i, (box, prob) in enumerate(zip(detection[0], detection[1])):
+        if prob < 0.9:  # Filter out detections with low confidence
+            continue
+        if (box[2] - box[0] < min_size[0]) or (box[3] - box[1] < min_size[1]):  # Check size
+            continue
+        return box  # Return the first valid detection
+    return None
+
 @spaces.GPU
 def process_frame(frame):
-    face = mtcnn(frame)
-    if face is None:
+    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    detection = mtcnn.detect(rgb_frame)
+    face_box = filter_low_quality_detections(detection)
+    
+    if face_box is None:
         return None, None, None
 
-    face = face.unsqueeze(0)
-    face = F.interpolate(face, size=(256, 256), mode='bilinear', align_corners=False)
+    x1, y1, x2, y2 = map(int, face_box)
+    h, w, _ = rgb_frame.shape
+
+    if x1 < 0 or y1 < 0 or x2 > w or y2 > h:
+        return None, None, None
+
+    face = rgb_frame[y1:y2, x1:x2]
+    if face.size == 0:
+        return None, None, None
 
-    face = face.to(DEVICE)
-    face = face.to(torch.float32)
-    face = face / 255.0
+    face = cv2.resize(face, (256, 256))
+    face = torch.from_numpy(face).permute(2, 0, 1).unsqueeze(0).to(DEVICE)
+    face = face.to(torch.float32) / 255.0
 
     with torch.no_grad():
         output = torch.sigmoid(model(face).squeeze(0))
@@ -82,8 +106,7 @@ def analyze_video(video_path, sample_rate=30, top_n=5, detection_threshold=0.5):
             break
 
         if frame_count % sample_rate == 0:
-            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            prediction, confidence, visualization = process_frame(rgb_frame)
+            prediction, confidence, visualization = process_frame(frame)
 
             if prediction is not None:
                 total_processed += 1