Update app.py

app.py

@@ -239,12 +239,19 @@ class ForgeryDetector:
             logits, decoder_features = self.model(image_tensor)
             prob_map = torch.sigmoid(logits).cpu().numpy()[0, 0]
         
+        # Resize probability map to match original image size to avoid index mismatch errors
+        prob_map_resized = cv2.resize(
+            prob_map, 
+            (original_image.shape[1], original_image.shape[0]), 
+            interpolation=cv2.INTER_LINEAR
+        )
+        
         # Refine mask
-        binary_mask = (prob_map > 0.5).astype(np.uint8)
-        refined_mask = self.mask_refiner.refine(binary_mask, original_size=original_image.shape[:2])
+        binary_mask = (prob_map_resized > 0.5).astype(np.uint8)
+        refined_mask = self.mask_refiner.refine(prob_map_resized, original_size=original_image.shape[:2])
         
         # Extract regions
-        regions = self.region_extractor.extract(refined_mask, prob_map, original_image)
+        regions = self.region_extractor.extract(refined_mask, prob_map_resized, original_image)
         
         # Classify regions
         results = []
@@ -461,7 +468,7 @@ with gr.Blocks(css=custom_css) as demo:
             input_file = gr.Image(
                 label="Document (Image or PDF)",
                 type="filepath",
-                sources=["upload"]
+                sources=["upload", "webcam"]
             )
             
             with gr.Row():