deal with large images

app.py

@@ -5,6 +5,7 @@ import sys
 import json
 import argparse
 from types import SimpleNamespace
+from PIL import Image
 
 # Try to import detector - if this fails, we'll show an error in the UI
 try:
@@ -30,6 +31,42 @@ if os.environ.get("SPACE_ID"):
 # Available detectors based on launcher.py
 DETECTORS = ['R50_TF', 'R50_nodown', 'CLIP-D', 'P2G', 'NPR']
 
+def process_image(image_path):
+    """
+    Check if image is larger than 1024x1024 and central crop it if necessary.
+    Returns the path to the processed image (or original if no change).
+    """
+    try:
+        with Image.open(image_path) as img:
+            width, height = img.size
+            
+            # Check if both dimensions are larger than 1024
+            if width > 1024 and height > 1024:
+                print(f"Image size {width}x{height} exceeds 1024x1024. Performing central crop.")
+                
+                # Calculate crop box
+                left = (width - 1024) / 2
+                top = (height - 1024) / 2
+                right = (width + 1024) / 2
+                bottom = (height + 1024) / 2
+                
+                # Crop
+                img_cropped = img.crop((left, top, right, bottom))
+                
+                # Save to new path
+                directory, filename = os.path.split(image_path)
+                name, ext = os.path.splitext(filename)
+                new_filename = f"{name}_cropped{ext}"
+                new_path = os.path.join(directory, new_filename)
+                
+                img_cropped.save(new_path)
+                return new_path
+                
+            return image_path
+    except Exception as e:
+        print(f"Error processing image: {e}")
+        return image_path
+
 def predict(image_path, detector_name):
     # Check if detector is available
     if not DETECTOR_AVAILABLE:
@@ -42,12 +79,20 @@ def predict(image_path, detector_name):
     if not image_path:
         return json.dumps({"error": "Please upload an image."}, indent=2)
     
+    # Process image (central crop if too large)
+    processed_path = image_path
+    try:
+        processed_path = process_image(image_path)
+    except Exception as e:
+        print(f"Warning: Image processing failed: {e}")
+        # Continue with original image if processing fails
+    
     # Create a temporary output file path
     output_path = "temp_result.json"
     
     # Mock args object
     args = SimpleNamespace(
-        image=image_path,
+        image=processed_path,
         detector=detector_name,
         config_dir='configs',
         output=output_path,
@@ -99,6 +144,13 @@ def predict(image_path, detector_name):
         # Cleanup
         if os.path.exists(output_path):
             os.remove(output_path)
+        
+        # Cleanup cropped image if it's different from original
+        if processed_path != image_path and os.path.exists(processed_path):
+            try:
+                os.remove(processed_path)
+            except Exception as e:
+                print(f"Warning: Could not remove temporary file {processed_path}: {e}")
 
 # Create Gradio Interface
 # Use theme only if gradio version supports it
@@ -119,7 +171,7 @@ with demo:
     *   **Elapsed Time**: The time the model needed to make the prediction (excluding preprocessing or model building).
 
     ### Note
-    ⚠️ Due to file size limitations, model weights need to be downloaded automatically on first use. This may take a few moments.
+    ⚠️ Due to file size limitations, model weights need to be downloaded automatically on first use. This may take a few moments. \
     ⚠️ To provide a free service, all models run on CPU.The detection process may take a few seconds, depending on the image size and the selected detector.
     """)
     

test_cleanup.py

@@ -0,0 +1,56 @@
+import os
+import json
+from PIL import Image
+from app import predict, process_image
+
+# Mock DETECTOR_AVAILABLE to avoid import errors during test
+import app
+app.DETECTOR_AVAILABLE = True
+app.run_detect = lambda args: None # Mock run_detect
+
+def test_cleanup():
+    # Create a large dummy image
+    large_img_path = "test_large_cleanup.png"
+    Image.new('RGB', (2000, 2000), color='green').save(large_img_path)
+    
+    # Create a dummy output file because predict expects it
+    output_path = "temp_result.json"
+    with open(output_path, 'w') as f:
+        json.dump({"prediction": "real", "confidence": 0.9, "elapsed_time": 0.1}, f)
+        
+    # We need to monkeypatch process_image to track the file it creates
+    original_process_image = app.process_image
+    created_files = []
+    
+    def tracked_process_image(path):
+        new_path = original_process_image(path)
+        if new_path != path:
+            created_files.append(new_path)
+        return new_path
+        
+    app.process_image = tracked_process_image
+    
+    # Run predict
+    try:
+        print("Running predict...")
+        app.predict(large_img_path, "R50_TF")
+    except Exception as e:
+        print(f"Predict failed: {e}")
+        
+    # Check if cropped file was created
+    assert len(created_files) > 0, "Should have created a cropped file"
+    cropped_path = created_files[0]
+    print(f"Cropped path was: {cropped_path}")
+    
+    # Check if cropped file was deleted
+    assert not os.path.exists(cropped_path), f"Cropped file {cropped_path} should have been deleted"
+    print("Cleanup test passed!")
+    
+    # Cleanup original
+    if os.path.exists(large_img_path):
+        os.remove(large_img_path)
+    if os.path.exists(output_path):
+        os.remove(output_path)
+
+if __name__ == "__main__":
+    test_cleanup()

test_crop.py

@@ -0,0 +1,44 @@
+import os
+from PIL import Image
+from app import process_image
+
+def test_process_image():
+    # Create a large dummy image
+    large_img_path = "test_large.png"
+    Image.new('RGB', (2000, 2000), color='red').save(large_img_path)
+    
+    # Process it
+    processed_path = process_image(large_img_path)
+    
+    # Check if a new file was created
+    assert processed_path != large_img_path, "Should have returned a new path"
+    assert "cropped" in processed_path, "New path should contain 'cropped'"
+    
+    # Check dimensions
+    with Image.open(processed_path) as img:
+        assert img.size == (1024, 1024), f"Expected 1024x1024, got {img.size}"
+        
+    print("Large image test passed!")
+    
+    # Create a small dummy image
+    small_img_path = "test_small.png"
+    Image.new('RGB', (800, 800), color='blue').save(small_img_path)
+    
+    # Process it
+    processed_path_small = process_image(small_img_path)
+    
+    # Check if it returned the same path
+    assert processed_path_small == small_img_path, "Should have returned original path"
+    
+    print("Small image test passed!")
+    
+    # Cleanup
+    if os.path.exists(large_img_path):
+        os.remove(large_img_path)
+    if os.path.exists(processed_path):
+        os.remove(processed_path)
+    if os.path.exists(small_img_path):
+        os.remove(small_img_path)
+
+if __name__ == "__main__":
+    test_process_image()