""" Enhanced File Upload Handler for CAD2Program Handles various file types, batch uploads, and file validation """ import os import tempfile import zipfile import tarfile import json import shutil from pathlib import Path from typing import Dict, List, Tuple, Optional, Any, Union import logging from PIL import Image, ImageOps import io import base64 import mimetypes logger = logging.getLogger(__name__) class FileUploadHandler: """Enhanced file upload handler with validation and processing""" def __init__(self, upload_dir: str = "uploads", max_file_size: int = 50 * 1024 * 1024): self.upload_dir = Path(upload_dir) self.upload_dir.mkdir(exist_ok=True) self.max_file_size = max_file_size # 50MB default # Supported file types self.supported_image_types = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.webp'} self.supported_program_types = {'.py', '.txt'} self.supported_archive_types = {'.zip', '.tar', '.tar.gz', '.tgz'} self.supported_cad_types = {'.dxf', '.dwg', '.step', '.stp', '.stl', '.obj'} def validate_file(self, file_path: Union[str, Path]) -> Dict[str, Any]: """Validate uploaded file""" file_path = Path(file_path) validation_result = { "valid": False, "file_type": None, "size": 0, "errors": [], "warnings": [] } try: # Check if file exists if not file_path.exists(): validation_result["errors"].append("File does not exist") return validation_result # Check file size file_size = file_path.stat().st_size validation_result["size"] = file_size if file_size > self.max_file_size: validation_result["errors"].append(f"File too large: {file_size / 1024 / 1024:.1f}MB > {self.max_file_size / 1024 / 1024:.1f}MB") return validation_result if file_size == 0: validation_result["errors"].append("File is empty") return validation_result # Check file extension file_extension = file_path.suffix.lower() if file_extension in self.supported_image_types: validation_result["file_type"] = "image" # Validate image try: with Image.open(file_path) as img: validation_result["image_info"] = { "format": img.format, "mode": img.mode, "size": img.size } # Check image dimensions if img.size[0] < 64 or img.size[1] < 64: validation_result["warnings"].append("Image resolution is very low") elif img.size[0] > 4096 or img.size[1] > 4096: validation_result["warnings"].append("Image resolution is very high, will be resized") except Exception as e: validation_result["errors"].append(f"Invalid image file: {str(e)}") return validation_result elif file_extension in self.supported_program_types: validation_result["file_type"] = "program" # Validate program file try: with open(file_path, 'r', encoding='utf-8') as f: content = f.read() validation_result["program_info"] = { "lines": len(content.split('\n')), "characters": len(content), "has_primitives": "add_primitive" in content or "Primitive" in content } except Exception as e: validation_result["errors"].append(f"Cannot read program file: {str(e)}") return validation_result elif file_extension in self.supported_archive_types: validation_result["file_type"] = "archive" # Validate archive try: if file_extension == '.zip': with zipfile.ZipFile(file_path, 'r') as zf: validation_result["archive_info"] = { "format": "zip", "files": len(zf.namelist()), "file_list": zf.namelist()[:10] # First 10 files } elif file_extension in {'.tar', '.tar.gz', '.tgz'}: with tarfile.open(file_path, 'r') as tf: validation_result["archive_info"] = { "format": "tar", "files": len(tf.getnames()), "file_list": tf.getnames()[:10] # First 10 files } except Exception as e: validation_result["errors"].append(f"Invalid archive file: {str(e)}") return validation_result elif file_extension in self.supported_cad_types: validation_result["file_type"] = "cad" validation_result["cad_info"] = { "format": file_extension[1:].upper(), "note": "CAD file detected - may need special processing" } else: validation_result["errors"].append(f"Unsupported file type: {file_extension}") return validation_result # If we get here, file is valid validation_result["valid"] = True except Exception as e: validation_result["errors"].append(f"Validation error: {str(e)}") return validation_result def process_single_file(self, file_data: Any, filename: str) -> Dict[str, Any]: """Process a single uploaded file""" # Create unique filename to avoid conflicts timestamp = str(int(os.path.getmtime(filename)) if os.path.exists(filename) else 0) safe_filename = self._sanitize_filename(filename) unique_filename = f"{timestamp}_{safe_filename}" file_path = self.upload_dir / unique_filename try: # Save file if hasattr(file_data, 'read'): # File-like object with open(file_path, 'wb') as f: f.write(file_data.read()) elif isinstance(file_data, bytes): # Bytes data with open(file_path, 'wb') as f: f.write(file_data) else: # Assume it's a path shutil.copy2(file_data, file_path) # Validate file validation = self.validate_file(file_path) if not validation["valid"]: # Clean up invalid file if file_path.exists(): file_path.unlink() return { "success": False, "filename": filename, "errors": validation["errors"], "validation": validation } # Process based on file type processed_data = self._process_by_type(file_path, validation) return { "success": True, "filename": filename, "saved_path": str(file_path), "validation": validation, "processed_data": processed_data } except Exception as e: logger.error(f"Error processing file {filename}: {e}") return { "success": False, "filename": filename, "errors": [f"Processing error: {str(e)}"] } def process_batch_upload(self, files: List[Tuple[Any, str]]) -> Dict[str, Any]: """Process multiple files at once""" results = { "successful": [], "failed": [], "summary": { "total": len(files), "success_count": 0, "error_count": 0, "images": 0, "programs": 0, "archives": 0 } } for file_data, filename in files: result = self.process_single_file(file_data, filename) if result["success"]: results["successful"].append(result) results["summary"]["success_count"] += 1 # Count file types file_type = result["validation"]["file_type"] if file_type in results["summary"]: results["summary"][file_type + "s"] += 1 else: results["failed"].append(result) results["summary"]["error_count"] += 1 return results def process_archive(self, archive_path: Union[str, Path]) -> Dict[str, Any]: """Extract and process archive contents""" archive_path = Path(archive_path) extract_dir = self.upload_dir / f"extracted_{archive_path.stem}" extract_dir.mkdir(exist_ok=True) extracted_files = [] try: # Extract archive if archive_path.suffix.lower() == '.zip': with zipfile.ZipFile(archive_path, 'r') as zf: zf.extractall(extract_dir) extracted_files = [extract_dir / name for name in zf.namelist()] elif archive_path.suffix.lower() in {'.tar', '.tar.gz', '.tgz'}: with tarfile.open(archive_path, 'r') as tf: tf.extractall(extract_dir) extracted_files = [extract_dir / name for name in tf.getnames()] # Process extracted files processed_results = [] for extracted_file in extracted_files: if extracted_file.is_file(): validation = self.validate_file(extracted_file) if validation["valid"]: processed_data = self._process_by_type(extracted_file, validation) processed_results.append({ "file_path": str(extracted_file), "validation": validation, "processed_data": processed_data }) return { "success": True, "extract_dir": str(extract_dir), "extracted_count": len(extracted_files), "processed_count": len(processed_results), "results": processed_results } except Exception as e: return { "success": False, "error": f"Archive processing failed: {str(e)}" } def _process_by_type(self, file_path: Path, validation: Dict) -> Dict[str, Any]: """Process file based on its type""" file_type = validation["file_type"] processed_data = {"type": file_type} if file_type == "image": processed_data.update(self._process_image(file_path, validation)) elif file_type == "program": processed_data.update(self._process_program(file_path, validation)) elif file_type == "archive": processed_data.update(self.process_archive(file_path)) elif file_type == "cad": processed_data.update(self._process_cad_file(file_path, validation)) return processed_data def _process_image(self, file_path: Path, validation: Dict) -> Dict[str, Any]: """Process uploaded image file""" try: with Image.open(file_path) as img: # Convert to RGB if necessary if img.mode != 'RGB': img = img.convert('RGB') # Resize if too large max_size = 1024 if img.size[0] > max_size or img.size[1] > max_size: img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS) # Save resized image resized_path = file_path.parent / f"resized_{file_path.name}" img.save(resized_path, 'PNG') return { "original_size": validation["image_info"]["size"], "resized_path": str(resized_path), "new_size": img.size, "format": "PNG" } return { "size": img.size, "format": img.format, "mode": img.mode, "processed": True } except Exception as e: return {"error": f"Image processing failed: {str(e)}"} def _process_program(self, file_path: Path, validation: Dict) -> Dict[str, Any]: """Process uploaded program file""" try: with open(file_path, 'r', encoding='utf-8') as f: content = f.read() # Analyze program content analysis = { "lines": len(content.split('\n')), "characters": len(content), "primitives": [], "imports": [], "functions": [] } # Extract primitives for line in content.split('\n'): line = line.strip() if 'add_primitive' in line: analysis["primitives"].append(line) elif line.startswith('import ') or line.startswith('from '): analysis["imports"].append(line) elif line.startswith('def '): analysis["functions"].append(line) return { "analysis": analysis, "content_preview": content[:500], # First 500 characters "processed": True } except Exception as e: return {"error": f"Program processing failed: {str(e)}"} def _process_cad_file(self, file_path: Path, validation: Dict) -> Dict[str, Any]: """Process CAD files (placeholder for future implementation)""" return { "format": validation["cad_info"]["format"], "note": "CAD file processing not yet implemented", "future_feature": True } def create_dataset_from_uploads(self, upload_results: Dict) -> Dict[str, Any]: """Create a dataset from processed uploads""" dataset_entries = [] for result in upload_results["successful"]: validation = result["validation"] processed = result["processed_data"] if validation["file_type"] == "image": # Look for matching program file image_path = Path(result["saved_path"]) base_name = image_path.stem # Try to find matching program program_path = None for prog_result in upload_results["successful"]: if prog_result["validation"]["file_type"] == "program": prog_path = Path(prog_result["saved_path"]) if prog_path.stem == base_name: program_path = prog_path break if program_path: # Read program content with open(program_path, 'r') as f: program_content = f.read() entry = { "id": f"upload_{base_name}", "image_path": str(image_path), "program_path": str(program_path), "program_text": program_content, "source": "user_upload", "validation": validation, "processed_data": processed } dataset_entries.append(entry) # Save dataset dataset_file = self.upload_dir / "created_dataset.json" with open(dataset_file, 'w') as f: json.dump(dataset_entries, f, indent=2) return { "dataset_file": str(dataset_file), "entries": len(dataset_entries), "pairs_found": len(dataset_entries) } def _sanitize_filename(self, filename: str) -> str: """Sanitize filename for safe storage""" # Remove path separators and dangerous characters safe_chars = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-") sanitized = "".join(c if c in safe_chars else "_" for c in filename) # Limit length if len(sanitized) > 100: name, ext = os.path.splitext(sanitized) sanitized = name[:100-len(ext)] + ext return sanitized def cleanup_old_uploads(self, days: int = 7): """Clean up old uploaded files""" import time cutoff_time = time.time() - (days * 24 * 60 * 60) cleaned_files = 0 for file_path in self.upload_dir.rglob('*'): if file_path.is_file(): try: if file_path.stat().st_mtime < cutoff_time: file_path.unlink() cleaned_files += 1 except Exception as e: logger.warning(f"Failed to clean up {file_path}: {e}") return {"cleaned_files": cleaned_files} def get_upload_stats(self) -> Dict[str, Any]: """Get statistics about uploads""" stats = { "total_files": 0, "total_size": 0, "file_types": {}, "recent_uploads": [] } for file_path in self.upload_dir.rglob('*'): if file_path.is_file(): stats["total_files"] += 1 file_size = file_path.stat().st_size stats["total_size"] += file_size # Count by extension ext = file_path.suffix.lower() stats["file_types"][ext] = stats["file_types"].get(ext, 0) + 1 # Track recent uploads (last 24 hours) import time if file_path.stat().st_mtime > time.time() - 86400: stats["recent_uploads"].append({ "name": file_path.name, "size": file_size, "modified": file_path.stat().st_mtime }) return stats # Gradio-specific upload handlers def create_gradio_upload_interface(): """Create Gradio interface components for file uploads""" import gradio as gr upload_handler = FileUploadHandler() def handle_gradio_upload(files): """Handle file upload from Gradio interface""" if not files: return "No files uploaded", "" # Process files file_tuples = [] for file in files: file_tuples.append((file.file, file.name)) results = upload_handler.process_batch_upload(file_tuples) # Format results for display summary = f""" ## Upload Results **Total Files:** {results['summary']['total']} **Successful:** {results['summary']['success_count']} **Failed:** {results['summary']['error_count']} **File Types:** - Images: {results['summary']['images']} - Programs: {results['summary']['programs']} - Archives: {results['summary']['archives']} """ # Create dataset if we have matching pairs dataset_info = "" if results['summary']['success_count'] > 1: dataset = upload_handler.create_dataset_from_uploads(results) dataset_info = f"\n**Dataset Created:** {dataset['entries']} entries" return summary + dataset_info, str(results) # Define Gradio components file_upload = gr.File( label="Upload CAD Files", file_count="multiple", file_types=[".png", ".jpg", ".jpeg", ".py", ".txt", ".zip"] ) upload_btn = gr.Button("Process Uploads", variant="primary") results_display = gr.Markdown(label="Upload Results") detailed_results = gr.JSON(label="Detailed Results", visible=False) upload_btn.click( fn=handle_gradio_upload, inputs=[file_upload], outputs=[results_display, detailed_results] ) return file_upload, upload_btn, results_display, detailed_results # Utility functions for integration with main app def setup_upload_system(base_dir: str = "uploads") -> FileUploadHandler: """Setup upload system with proper directory structure""" handler = FileUploadHandler(base_dir) # Create subdirectories (Path(base_dir) / "images").mkdir(exist_ok=True) (Path(base_dir) / "programs").mkdir(exist_ok=True) (Path(base_dir) / "archives").mkdir(exist_ok=True) (Path(base_dir) / "processed").mkdir(exist_ok=True) return handler if __name__ == "__main__": # Test upload handler handler = setup_upload_system("test_uploads") print("Upload handler created") print(f"Upload directory: {handler.upload_dir}") # Test file validation test_extensions = ['.png', '.py', '.zip', '.invalid'] for ext in test_extensions: # Create dummy file for testing test_file = handler.upload_dir / f"test{ext}" test_file.write_text("test content") validation = handler.validate_file(test_file) print(f"{ext}: Valid={validation['valid']}, Type={validation.get('file_type', 'unknown')}") # Clean up test_file.unlink() # Get stats stats = handler.get_upload_stats() print(f"Upload stats: {stats}")