Update app/main.py

app/main.py

@@ -72,19 +72,32 @@ async def predict(file: UploadFile = File(...)):
             if not success:
                 raise HTTPException(status_code=500, detail="Model failed to load")
         
-        # Validate file type
-        if not file.content_type.startswith('image/'):
+        # Validate file type - more robust approach
+        # Don't rely solely on content_type as it might be incorrect
+        valid_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.gif', '.webp']
+        filename_lower = (file.filename or '').lower()
+        
+        # Check both content type and file extension
+        is_valid_content_type = file.content_type and file.content_type.startswith('image/')
+        is_valid_extension = any(filename_lower.endswith(ext) for ext in valid_extensions)
+        
+        if not (is_valid_content_type or is_valid_extension):
+            logger.warning(f"Invalid file type: content_type={file.content_type}, filename={file.filename}")
             raise HTTPException(status_code=400, detail="File must be an image")
         
         # Read and process image
         logger.info(f"Processing image: {file.filename}")
         image_data = await file.read()
         
-        # Optimize image processing
+        # Optimize image processing with better error handling
         try:
             image = Image.open(io.BytesIO(image_data))
-            # Convert to RGB and resize if too large
-            image = image.convert("RGB")
+            
+            # Verify it's actually an image by trying to get basic info
+            image.verify()  # This will raise an exception if not a valid image
+            
+            # Reopen the image since verify() closes it
+            image = Image.open(io.BytesIO(image_data)).convert("RGB")
             
             # Resize large images to improve speed
             max_size = (1024, 1024)
@@ -94,7 +107,7 @@ async def predict(file: UploadFile = File(...)):
             
         except Exception as e:
             logger.error(f"Image processing error: {e}")
-            raise HTTPException(status_code=400, detail="Invalid image format")
+            raise HTTPException(status_code=400, detail="Invalid or corrupted image file")
         
         # Run prediction with timeout
         try: