{ "cells": [ { "cell_type": "code", "execution_count": 3, "id": "0416e9bd", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Current directory: C:\\Users\\HP\\NitroSense-AI\n", "โœ… Created U2Net folders\n", "โœ… U2Net weights found: C:\\Users\\HP\\NitroSense-AI\\u2net\\weights\\u2net.pth\n", " File size: 168.12 MB\n", "\n", "======================================================================\n", "๐Ÿงช TESTING PORTABLE BACKGROUND REMOVER WITH ACTUAL WEIGHTS\n", "======================================================================\n", "\n", "โœ… Found weights at: C:\\Users\\HP\\NitroSense-AI\\u2net\\weights\\u2net.pth\n", "\n", "๐Ÿ”„ Initializing BackgroundRemover...\n", "Using device: cpu\n", "โœ… U2Net source already present\n", "โœ… U2Net module imported\n", "๐Ÿ“ฆ Loading U2Net model from: C:\\Users\\HP\\NitroSense-AI\\u2net\\weights\\u2net.pth\n", "โœ… U2Net model loaded successfully!\n", " Model type: U2NET\n", "โœ… BackgroundRemover initialized with U2Net model\n", "\n", "๐Ÿ“ท Testing with: 2888723a-fd6f-4471-b7e5-692727bd5692_668bc6dbe4b48.jpg\n", " Original shape: (1717, 3162, 3)\n", " Original dtype: uint8\n", " Original range: [0, 241]\n", "\n", "๐Ÿ”„ Test 1: Basic background removal...\n", " ๐ŸŽจ Background removal - Original colors preserved\n", " Mask range: [0, 254]\n", " Background pixels: 38454/50176\n", " โœ… Result shape: (224, 224, 3)\n", " โœ… Result dtype: uint8\n", " โœ… Result range: [0, 255]\n", " โœ… Mask shape: (224, 224)\n", " โœ… Mask range: [0, 254]\n", "\n", "๐Ÿ”„ Test 2: Process and save...\n", " ๐ŸŽจ Background removal - Original colors preserved\n", " Mask range: [0, 254]\n", " Background pixels: 38454/50176\n", "โœ… Saved results to C:\\Users\\HP\\NitroSense-AI\\test_output_with_weights\n", " โœ… Files saved to: C:\\Users\\HP\\NitroSense-AI\\test_output_with_weights\n", " - original: 2888723a-fd6f-4471-b7e5-692727bd5692_668bc6dbe4b48_original.jpg (20.1 KB)\n", " - mask: 2888723a-fd6f-4471-b7e5-692727bd5692_668bc6dbe4b48_mask.jpg (10.2 KB)\n", " - result: 2888723a-fd6f-4471-b7e5-692727bd5692_668bc6dbe4b48_bg_removed.jpg (15.9 KB)\n", "\n", "โœ… All tests completed!\n", " Mode: FULL\n", "\n", "๐Ÿ“ Check the saved images in: C:\\Users\\HP\\NitroSense-AI\\test_output_with_weights\n", "\n", "โœ… Saved portable BackgroundRemover to C:\\Users\\HP\\NitroSense-AI\\utils\\background_removal.py\n", "โœ… Created setup_u2net.py - run this once before deploying\n", "โœ… Created requirements.txt\n", "\n", "======================================================================\n", "๐ŸŽฏ PORTABLE BACKGROUND REMOVER - FINAL CHECK\n", "======================================================================\n", "\n", "โœ… Created files:\n", " โ€ข utils/background_removal.py - Portable U2Net class\n", " โ€ข setup_u2net.py - Run once to download U2Net source\n", " โ€ข requirements.txt - All dependencies\n", "\n", "๐Ÿ“ Current folder structure:\n", " NitroSense-AI/\n", " โ”œโ”€โ”€ app.py\n", " โ”œโ”€โ”€ requirements.txt\n", " โ”œโ”€โ”€ setup_u2net.py\n", " โ”œโ”€โ”€ u2net/\n", " โ”‚ โ”œโ”€โ”€ model/\n", " โ”‚ โ”‚ โ””โ”€โ”€ u2net.py (source code)\n", " โ”‚ โ””โ”€โ”€ weights/\n", " โ”‚ โ””โ”€โ”€ u2net.pth โœ… (weights)\n", " โ””โ”€โ”€ utils/\n", " โ””โ”€โ”€ background_removal.py\n", "\n", "โœ… U2Net weights are present! Background removal will work.\n", "\n", "๐Ÿ“ Test results saved to: C:\\Users\\HP\\NitroSense-AI\\test_output_with_weights\n", " Check the following files:\n", " - 6538a45fe4f02_original.jpg (original resized image)\n", " - 6538a45fe4f02_mask.jpg (U2Net mask)\n", " - 6538a45fe4f02_bg_removed.jpg (final result)\n", "\n", "๐Ÿ“‹ Deployment steps:\n", "1. Run: python setup_u2net.py (downloads U2Net source)\n", "2. Ensure u2net.pth is in u2net/weights/\n", "3. Install: pip install -r requirements.txt\n", "4. Run: python app.py\n", "\n", "โœ… Your app is now fully portable!\n", " Works on any system - no hardcoded paths!\n" ] } ], "source": [ "# %% [markdown]\n", "# # Step 2: Portable U2Net Background Removal (Memory Optimized - No Display)\n", "# ## Same training pipeline but with maximum memory efficiency and portability\n", "\n", "# %%\n", "import torch\n", "import numpy as np\n", "import cv2\n", "from torchvision import transforms\n", "from PIL import Image\n", "import sys\n", "import os\n", "import gc\n", "import urllib.request\n", "import zipfile\n", "import shutil\n", "import warnings\n", "warnings.filterwarnings('ignore')\n", "\n", "# Force CPU and aggressive memory settings\n", "torch.set_num_threads(1)\n", "os.environ[\"OMP_NUM_THREADS\"] = \"1\"\n", "os.environ[\"MKL_NUM_THREADS\"] = \"1\"\n", "\n", "# %%\n", "# Get current directory\n", "current_dir = os.getcwd()\n", "print(f\"Current directory: {current_dir}\")\n", "\n", "# Create u2net folder structure\n", "u2net_folder = os.path.join(current_dir, 'u2net')\n", "model_folder = os.path.join(u2net_folder, 'model')\n", "weights_folder = os.path.join(u2net_folder, 'weights')\n", "os.makedirs(model_folder, exist_ok=True)\n", "os.makedirs(weights_folder, exist_ok=True)\n", "print(f\"โœ… Created U2Net folders\")\n", "\n", "# %% [markdown]\n", "# ### Check if weights are downloaded\n", "\n", "# %%\n", "weights_path = os.path.join(weights_folder, 'u2net.pth')\n", "if os.path.exists(weights_path):\n", " file_size = os.path.getsize(weights_path) / (1024 * 1024) # MB\n", " print(f\"โœ… U2Net weights found: {weights_path}\")\n", " print(f\" File size: {file_size:.2f} MB\")\n", "else:\n", " print(f\"โŒ U2Net weights not found at: {weights_path}\")\n", " print(\" Please download from: https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ\")\n", "\n", "# %% [markdown]\n", "# ### Helper function to get a test image from the project\n", "\n", "# %%\n", "def get_test_image():\n", " \"\"\"Get the first image from static/uploads folder or create a dummy image\"\"\"\n", " # First, check static/uploads folder for existing images\n", " upload_folder = os.path.join(current_dir, 'static', 'uploads')\n", " if os.path.exists(upload_folder):\n", " images = [f for f in os.listdir(upload_folder) \n", " if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif'))]\n", " if images:\n", " return os.path.join(upload_folder, images[0])\n", " \n", " # If no images found, check if there's a sample image in test_output\n", " test_output = os.path.join(current_dir, 'test_output')\n", " if os.path.exists(test_output):\n", " images = [f for f in os.listdir(test_output) \n", " if f.lower().endswith(('.png', '.jpg', '.jpeg'))]\n", " if images:\n", " return os.path.join(test_output, images[0])\n", " \n", " # If no images found, return None (will use dummy image)\n", " return None\n", "\n", "# %% [markdown]\n", "# ### Portable U2Net Background Remover Class\n", "\n", "# %%\n", "class BackgroundRemover:\n", " \"\"\"\n", " Portable U2Net Background Remover\n", " Automatically downloads U2Net if not found\n", " Memory-optimized - no display, just processing\n", " \"\"\"\n", " \n", " def __init__(self, model_path=None):\n", " self.device = torch.device('cpu')\n", " print(f\"Using device: {self.device}\")\n", " \n", " # Get paths\n", " self.current_dir = os.path.dirname(os.path.abspath(__file__)) if '__file__' in dir() else os.getcwd()\n", " self.project_root = os.path.dirname(self.current_dir) if '__file__' in dir() else self.current_dir\n", " \n", " # Set up paths\n", " self.u2net_dir = os.path.join(self.project_root, 'u2net')\n", " self.model_dir = os.path.join(self.u2net_dir, 'model')\n", " self.weights_dir = os.path.join(self.u2net_dir, 'weights')\n", " \n", " # Create directories\n", " os.makedirs(self.model_dir, exist_ok=True)\n", " os.makedirs(self.weights_dir, exist_ok=True)\n", " \n", " # Download/verify U2Net source\n", " self.u2net_available = self._ensure_u2net_source()\n", " \n", " # Set model path\n", " if model_path is None:\n", " model_path = os.path.join(self.weights_dir, 'u2net.pth')\n", " self.model_path = model_path\n", " \n", " # Import and load model\n", " self._import_u2net()\n", " self._load_model()\n", " \n", " # Training transform\n", " self.transform = transforms.Compose([\n", " transforms.Resize((320, 320)),\n", " transforms.ToTensor(),\n", " transforms.Normalize(mean=[0.485, 0.456, 0.406],\n", " std=[0.229, 0.224, 0.225]),\n", " ])\n", " \n", " gc.collect()\n", " if self.model is not None:\n", " print(\"โœ… BackgroundRemover initialized with U2Net model\")\n", " else:\n", " print(\"โœ… BackgroundRemover initialized in fallback mode\")\n", " \n", " def _ensure_u2net_source(self):\n", " \"\"\"Download U2Net source if not present\"\"\"\n", " u2net_py = os.path.join(self.model_dir, 'u2net.py')\n", " \n", " if not os.path.exists(u2net_py):\n", " print(\"\\n๐Ÿ“ฅ Downloading U2Net source...\")\n", " u2net_url = \"https://github.com/xuebinqin/U-2-Net/archive/master.zip\"\n", " zip_path = os.path.join(self.u2net_dir, 'u2net.zip')\n", " \n", " try:\n", " # Download\n", " urllib.request.urlretrieve(u2net_url, zip_path)\n", " print(\" โœ… Downloaded\")\n", " \n", " # Extract\n", " print(\" ๐Ÿ“ฆ Extracting...\")\n", " with zipfile.ZipFile(zip_path, 'r') as zip_ref:\n", " zip_ref.extractall(self.u2net_dir)\n", " \n", " # Move model files\n", " extracted = os.path.join(self.u2net_dir, 'U-2-Net-master')\n", " if os.path.exists(extracted):\n", " src_model = os.path.join(extracted, 'model')\n", " if os.path.exists(src_model):\n", " for file in os.listdir(src_model):\n", " if file.endswith('.py'):\n", " shutil.copy(\n", " os.path.join(src_model, file),\n", " os.path.join(self.model_dir, file)\n", " )\n", " print(\" โœ… Model files copied\")\n", " \n", " # Clean up\n", " os.remove(zip_path)\n", " if os.path.exists(extracted):\n", " shutil.rmtree(extracted)\n", " \n", " print(\"โœ… U2Net source downloaded successfully\")\n", " return True\n", " \n", " except Exception as e:\n", " print(f\"โš ๏ธ Could not download U2Net: {e}\")\n", " print(\" Please manually download from:\")\n", " print(\" https://github.com/xuebinqin/U-2-Net\")\n", " print(f\" and place u2net.py in: {self.model_dir}\")\n", " return False\n", " else:\n", " print(\"โœ… U2Net source already present\")\n", " return True\n", " \n", " def _import_u2net(self):\n", " \"\"\"Import U2Net module\"\"\"\n", " if self.model_dir not in sys.path:\n", " sys.path.insert(0, self.model_dir)\n", " \n", " try:\n", " from u2net import U2NET\n", " self.U2NET = U2NET\n", " self.u2net_available = True\n", " print(\"โœ… U2Net module imported\")\n", " except ImportError as e:\n", " print(f\"โš ๏ธ Could not import U2Net: {e}\")\n", " print(\" Will use fallback mode\")\n", " self.U2NET = None\n", " self.u2net_available = False\n", " \n", " def _load_model(self):\n", " \"\"\"Load the U2Net model\"\"\"\n", " if self.U2NET is None:\n", " self.model = None\n", " print(\"โš ๏ธ Running in fallback mode (U2Net module missing)\")\n", " return\n", " \n", " if not os.path.exists(self.model_path):\n", " print(f\"โš ๏ธ Model weights not found at: {self.model_path}\")\n", " print(\" Will use fallback mode\")\n", " self.model = None\n", " return\n", " \n", " try:\n", " print(f\"๐Ÿ“ฆ Loading U2Net model from: {self.model_path}\")\n", " self.model = self.U2NET(in_ch=3, out_ch=1)\n", " state_dict = torch.load(\n", " self.model_path, \n", " map_location=self.device,\n", " weights_only=True\n", " )\n", " self.model.load_state_dict(state_dict)\n", " self.model.to(self.device)\n", " self.model.eval()\n", " print(\"โœ… U2Net model loaded successfully!\")\n", " print(f\" Model type: {type(self.model).__name__}\")\n", " except Exception as e:\n", " print(f\"โš ๏ธ Could not load model: {e}\")\n", " self.model = None\n", " \n", " def remove_background(self, image, target_size=(224, 224), max_size=800):\n", " \"\"\"\n", " Memory-optimized background removal - preserves original leaf colors\n", " \n", " Args:\n", " image: numpy array or PIL Image\n", " target_size: final size for model input\n", " max_size: maximum dimension for processing\n", " \n", " Returns:\n", " result: numpy array with background removed (original colors preserved)\n", " mask: binary mask from U2Net\n", " \"\"\"\n", " # Fallback if model not loaded\n", " if self.model is None:\n", " if isinstance(image, np.ndarray):\n", " return cv2.resize(image, target_size), None\n", " elif isinstance(image, Image.Image):\n", " return np.array(image.resize(target_size)), None\n", " return image, None\n", " \n", " try:\n", " # Store original image for color preservation\n", " original_image = image.copy() if isinstance(image, np.ndarray) else image\n", " \n", " # Resize large images for processing\n", " if isinstance(image, np.ndarray):\n", " h, w = image.shape[:2]\n", " if max(h, w) > max_size:\n", " scale = max_size / max(h, w)\n", " new_w = int(w * scale)\n", " new_h = int(h * scale)\n", " image = cv2.resize(image, (new_w, new_h))\n", " \n", " # Convert to PIL for U2Net\n", " image_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))\n", " elif isinstance(image, Image.Image):\n", " if max(image.size) > max_size:\n", " scale = max_size / max(image.size)\n", " new_size = tuple(int(dim * scale) for dim in image.size)\n", " image_pil = image.resize(new_size, Image.Resampling.LANCZOS)\n", " else:\n", " image_pil = image\n", " \n", " # Get mask from U2Net\n", " input_tensor = self.transform(image_pil).unsqueeze(0).to(self.device)\n", " \n", " with torch.no_grad():\n", " d1, d2, d3, d4, d5, d6, d7 = self.model(input_tensor)\n", " \n", " mask = d1[:, 0, :, :].cpu().numpy()\n", " \n", " # Clean up\n", " del d1, d2, d3, d4, d5, d6, d7, input_tensor\n", " gc.collect()\n", " \n", " # Process mask\n", " mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)\n", " mask = mask.transpose(1, 2, 0)\n", " \n", " # Resize mask to target size\n", " mask_target = cv2.resize(mask, target_size)\n", " mask_target = (mask_target * 255).astype(np.uint8)\n", " \n", " # === FIX: Preserve original colors ===\n", " # Get the original image resized to target size (without any color changes)\n", " if isinstance(original_image, np.ndarray):\n", " # Resize original image to target size\n", " image_resized = cv2.resize(original_image, target_size)\n", " # Ensure it's RGB\n", " if len(image_resized.shape) == 3 and image_resized.shape[2] == 3:\n", " # Keep as is - it's already in correct color space\n", " pass\n", " else:\n", " # It was a PIL Image\n", " image_resized = np.array(original_image.resize(target_size, Image.Resampling.LANCZOS))\n", " \n", " # Create result by copying the resized original\n", " result = image_resized.copy()\n", " \n", " # Apply mask - set background pixels to white, keep leaf pixels original\n", " # For RGB images, we need to handle 3-channel mask properly\n", " if len(mask_target.shape) == 3 and mask_target.shape[2] == 1:\n", " mask_target = mask_target.squeeze()\n", " \n", " # Set background (mask < threshold) to white\n", " result[mask_target < 128] = [255, 255, 255]\n", " \n", " # Debug info\n", " print(f\" ๐ŸŽจ Background removal - Original colors preserved\")\n", " print(f\" Mask range: [{mask_target.min()}, {mask_target.max()}]\")\n", " print(f\" Background pixels: {np.sum(mask_target < 128)}/{mask_target.size}\")\n", " \n", " return result, mask_target\n", " \n", " except Exception as e:\n", " print(f\"โš ๏ธ Warning in background removal: {e}\")\n", " import traceback\n", " traceback.print_exc()\n", " # Fallback: return resized original\n", " if isinstance(image, Image.Image):\n", " return np.array(image.resize(target_size)), None\n", " elif isinstance(image, np.ndarray):\n", " return cv2.resize(image, target_size), None\n", " return image, None\n", " \n", " def process_and_save(self, image_path, output_dir, target_size=(224, 224)):\n", " \"\"\"\n", " Process an image and save results without displaying\n", " \"\"\"\n", " try:\n", " # Read image\n", " img = cv2.imread(image_path)\n", " if img is None:\n", " print(f\"โŒ Could not read image: {image_path}\")\n", " return None\n", " \n", " img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n", " \n", " # Remove background\n", " result, mask = self.remove_background(img_rgb, target_size=target_size)\n", " \n", " # Save results\n", " os.makedirs(output_dir, exist_ok=True)\n", " \n", " base_name = os.path.splitext(os.path.basename(image_path))[0]\n", " \n", " # Save original resized\n", " orig_resized = cv2.resize(img_rgb, target_size)\n", " orig_resized_bgr = cv2.cvtColor(orig_resized, cv2.COLOR_RGB2BGR)\n", " cv2.imwrite(os.path.join(output_dir, f\"{base_name}_original.jpg\"), orig_resized_bgr)\n", " \n", " # Save mask if available\n", " if mask is not None:\n", " cv2.imwrite(os.path.join(output_dir, f\"{base_name}_mask.jpg\"), mask)\n", " \n", " # Save result\n", " result_bgr = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)\n", " cv2.imwrite(os.path.join(output_dir, f\"{base_name}_bg_removed.jpg\"), result_bgr)\n", " \n", " print(f\"โœ… Saved results to {output_dir}\")\n", " \n", " return {\n", " 'original': os.path.join(output_dir, f\"{base_name}_original.jpg\"),\n", " 'mask': os.path.join(output_dir, f\"{base_name}_mask.jpg\") if mask is not None else None,\n", " 'result': os.path.join(output_dir, f\"{base_name}_bg_removed.jpg\")\n", " }\n", " \n", " except Exception as e:\n", " print(f\"โŒ Error in process_and_save: {e}\")\n", " return None\n", " \n", " def __del__(self):\n", " \"\"\"Cleanup when object is deleted\"\"\"\n", " try:\n", " import gc\n", " gc.collect()\n", " except:\n", " pass\n", "\n", "# %% [markdown]\n", "# ### Test the Portable Background Remover with Actual Weights\n", "\n", "# %%\n", "print(\"\\n\" + \"=\"*70)\n", "print(\"๐Ÿงช TESTING PORTABLE BACKGROUND REMOVER WITH ACTUAL WEIGHTS\")\n", "print(\"=\"*70)\n", "\n", "# Check if weights exist\n", "if os.path.exists(weights_path):\n", " print(f\"\\nโœ… Found weights at: {weights_path}\")\n", " test_mode = \"full\"\n", "else:\n", " print(\"\\nโš ๏ธ U2Net weights not found!\")\n", " print(\" Testing in fallback mode (no actual background removal)\")\n", " test_mode = \"fallback\"\n", "\n", "# Initialize remover with weights\n", "print(\"\\n๐Ÿ”„ Initializing BackgroundRemover...\")\n", "remover = BackgroundRemover(weights_path if os.path.exists(weights_path) else None)\n", "\n", "# Get test image from project folders\n", "test_img = get_test_image()\n", "\n", "if test_img and os.path.exists(test_img):\n", " print(f\"\\n๐Ÿ“ท Testing with: {os.path.basename(test_img)}\")\n", " \n", " # Read image\n", " img = cv2.imread(test_img)\n", " img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n", " print(f\" Original shape: {img_rgb.shape}\")\n", " print(f\" Original dtype: {img_rgb.dtype}\")\n", " print(f\" Original range: [{img_rgb.min()}, {img_rgb.max()}]\")\n", " \n", " # Test 1: Basic removal\n", " print(\"\\n๐Ÿ”„ Test 1: Basic background removal...\")\n", " result, mask = remover.remove_background(img_rgb, target_size=(224, 224))\n", " \n", " print(f\" โœ… Result shape: {result.shape}\")\n", " print(f\" โœ… Result dtype: {result.dtype}\")\n", " print(f\" โœ… Result range: [{result.min()}, {result.max()}]\")\n", " \n", " if mask is not None:\n", " print(f\" โœ… Mask shape: {mask.shape}\")\n", " print(f\" โœ… Mask range: [{mask.min()}, {mask.max()}]\")\n", " else:\n", " print(f\" โ„น๏ธ No mask generated (fallback mode)\")\n", " \n", " # Test 2: Process and save\n", " print(\"\\n๐Ÿ”„ Test 2: Process and save...\")\n", " output_dir = os.path.join(current_dir, 'test_output_with_weights')\n", " saved = remover.process_and_save(test_img, output_dir)\n", " \n", " if saved:\n", " print(f\" โœ… Files saved to: {output_dir}\")\n", " for key, path in saved.items():\n", " if path:\n", " file_size = os.path.getsize(path) / 1024 # KB\n", " print(f\" - {key}: {os.path.basename(path)} ({file_size:.1f} KB)\")\n", " \n", " print(\"\\nโœ… All tests completed!\")\n", " print(f\" Mode: {test_mode.upper()}\")\n", " print(f\"\\n๐Ÿ“ Check the saved images in: {output_dir}\")\n", " \n", " # Clean up\n", " del remover\n", " gc.collect()\n", " \n", "else:\n", " print(f\"\\nโš ๏ธ No test image found in static/uploads/ or test_output/\")\n", " print(\" Creating dummy image for testing...\")\n", " \n", " # Create dummy image\n", " dummy = np.zeros((500, 500, 3), dtype=np.uint8)\n", " # Make it a green square on light blue background\n", " dummy[:, :, 0] = 100 # Blue component\n", " dummy[:, :, 1] = 150 # Green component\n", " dummy[:, :, 2] = 200 # Red component\n", " # Draw a green square in the center (simulating leaf)\n", " cv2.rectangle(dummy, (100, 100), (400, 400), (0, 255, 0), -1)\n", " # Add some texture inside the square\n", " cv2.rectangle(dummy, (150, 150), (350, 350), (0, 200, 0), -1)\n", " \n", " result, mask = remover.remove_background(dummy)\n", " print(f\"โœ… Dummy test passed - Result shape: {result.shape}\")\n", " if mask is not None:\n", " print(f\"โœ… Mask shape: {mask.shape}\")\n", " else:\n", " print(f\"โ„น๏ธ No mask generated (fallback mode)\")\n", "\n", "# %% [markdown]\n", "# ### Save the portable class for web app\n", "\n", "# %%\n", "# Save the portable class\n", "portable_code = '''\"\"\"\n", "Portable U2Net Background Remover\n", "Works on any system - downloads model if not present\n", "Memory-optimized for web applications\n", "\"\"\"\n", "import torch\n", "import numpy as np\n", "import cv2\n", "from torchvision import transforms\n", "from PIL import Image\n", "import sys\n", "import os\n", "import gc\n", "import urllib.request\n", "import zipfile\n", "import shutil\n", "import warnings\n", "warnings.filterwarnings('ignore')\n", "\n", "# Memory settings\n", "torch.set_num_threads(1)\n", "os.environ[\"OMP_NUM_THREADS\"] = \"1\"\n", "os.environ[\"MKL_NUM_THREADS\"] = \"1\"\n", "\n", "class BackgroundRemover:\n", " \"\"\"\n", " Portable U2Net Background Remover\n", " Automatically downloads U2Net if not found\n", " Memory-optimized for web applications\n", " \"\"\"\n", " \n", " def __init__(self, model_path=None):\n", " self.device = torch.device('cpu')\n", " \n", " # Get paths\n", " self.current_dir = os.path.dirname(os.path.abspath(__file__)) if '__file__' in dir() else os.getcwd()\n", " self.project_root = os.path.dirname(self.current_dir) if '__file__' in dir() else self.current_dir\n", " \n", " # Set up paths\n", " self.u2net_dir = os.path.join(self.project_root, 'u2net')\n", " self.model_dir = os.path.join(self.u2net_dir, 'model')\n", " self.weights_dir = os.path.join(self.u2net_dir, 'weights')\n", " \n", " # Create directories\n", " os.makedirs(self.model_dir, exist_ok=True)\n", " os.makedirs(self.weights_dir, exist_ok=True)\n", " \n", " # Download/verify U2Net source\n", " self.u2net_available = self._ensure_u2net_source()\n", " \n", " # Set model path\n", " if model_path is None:\n", " model_path = os.path.join(self.weights_dir, 'u2net.pth')\n", " self.model_path = model_path\n", " \n", " # Import and load model\n", " self._import_u2net()\n", " self._load_model()\n", " \n", " # Training transform\n", " self.transform = transforms.Compose([\n", " transforms.Resize((320, 320)),\n", " transforms.ToTensor(),\n", " transforms.Normalize(mean=[0.485, 0.456, 0.406],\n", " std=[0.229, 0.224, 0.225]),\n", " ])\n", " \n", " gc.collect()\n", " if self.model is not None:\n", " print(\"โœ… BackgroundRemover initialized with U2Net model\")\n", " else:\n", " print(\"โœ… BackgroundRemover initialized in fallback mode\")\n", " \n", " def _ensure_u2net_source(self):\n", " \"\"\"Download U2Net source if not present\"\"\"\n", " u2net_py = os.path.join(self.model_dir, 'u2net.py')\n", " \n", " if not os.path.exists(u2net_py):\n", " print(\"๐Ÿ“ฅ Downloading U2Net source...\")\n", " u2net_url = \"https://github.com/xuebinqin/U-2-Net/archive/master.zip\"\n", " zip_path = os.path.join(self.u2net_dir, 'u2net.zip')\n", " \n", " try:\n", " # Download\n", " urllib.request.urlretrieve(u2net_url, zip_path)\n", " \n", " # Extract\n", " with zipfile.ZipFile(zip_path, 'r') as zip_ref:\n", " zip_ref.extractall(self.u2net_dir)\n", " \n", " # Move model files\n", " extracted = os.path.join(self.u2net_dir, 'U-2-Net-master')\n", " if os.path.exists(extracted):\n", " src_model = os.path.join(extracted, 'model')\n", " if os.path.exists(src_model):\n", " for file in os.listdir(src_model):\n", " if file.endswith('.py'):\n", " shutil.copy(\n", " os.path.join(src_model, file),\n", " os.path.join(self.model_dir, file)\n", " )\n", " \n", " # Clean up\n", " os.remove(zip_path)\n", " if os.path.exists(extracted):\n", " shutil.rmtree(extracted)\n", " \n", " print(\"โœ… U2Net source downloaded successfully\")\n", " return True\n", " \n", " except Exception as e:\n", " print(f\"โš ๏ธ Could not download U2Net: {e}\")\n", " return False\n", " else:\n", " print(\"โœ… U2Net source already present\")\n", " return True\n", " \n", " def _import_u2net(self):\n", " \"\"\"Import U2Net module\"\"\"\n", " if self.model_dir not in sys.path:\n", " sys.path.insert(0, self.model_dir)\n", " \n", " try:\n", " from u2net import U2NET\n", " self.U2NET = U2NET\n", " self.u2net_available = True\n", " except ImportError:\n", " self.U2NET = None\n", " self.u2net_available = False\n", " \n", " def _load_model(self):\n", " \"\"\"Load the U2Net model\"\"\"\n", " if self.U2NET is None:\n", " self.model = None\n", " return\n", " \n", " if not os.path.exists(self.model_path):\n", " self.model = None\n", " return\n", " \n", " try:\n", " self.model = self.U2NET(in_ch=3, out_ch=1)\n", " state_dict = torch.load(\n", " self.model_path, \n", " map_location=self.device,\n", " weights_only=True\n", " )\n", " self.model.load_state_dict(state_dict)\n", " self.model.to(self.device)\n", " self.model.eval()\n", " except Exception:\n", " self.model = None\n", " \n", " def remove_background(self, image, target_size=(224, 224), max_size=800):\n", " \"\"\"\n", " Memory-optimized background removal - preserves original leaf colors\n", " \n", " Args:\n", " image: numpy array or PIL Image\n", " target_size: final size for model input\n", " max_size: maximum dimension for processing\n", " \n", " Returns:\n", " result: numpy array with background removed (original colors preserved)\n", " mask: binary mask from U2Net\n", " \"\"\"\n", " # Fallback if model not loaded\n", " if self.model is None:\n", " if isinstance(image, np.ndarray):\n", " return cv2.resize(image, target_size), None\n", " elif isinstance(image, Image.Image):\n", " return np.array(image.resize(target_size)), None\n", " return image, None\n", " \n", " try:\n", " # Store original image for color preservation\n", " original_image = image.copy() if isinstance(image, np.ndarray) else image\n", " \n", " # Resize large images for processing\n", " if isinstance(image, np.ndarray):\n", " h, w = image.shape[:2]\n", " if max(h, w) > max_size:\n", " scale = max_size / max(h, w)\n", " new_w = int(w * scale)\n", " new_h = int(h * scale)\n", " image = cv2.resize(image, (new_w, new_h))\n", " \n", " image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))\n", " elif isinstance(image, Image.Image):\n", " if max(image.size) > max_size:\n", " scale = max_size / max(image.size)\n", " new_size = tuple(int(dim * scale) for dim in image.size)\n", " image = image.resize(new_size, Image.Resampling.LANCZOS)\n", " else:\n", " image = image\n", " \n", " # Get mask\n", " input_tensor = self.transform(image).unsqueeze(0).to(self.device)\n", " \n", " with torch.no_grad():\n", " d1, d2, d3, d4, d5, d6, d7 = self.model(input_tensor)\n", " \n", " mask = d1[:, 0, :, :].cpu().numpy()\n", " \n", " # Clean up\n", " del d1, d2, d3, d4, d5, d6, d7, input_tensor\n", " gc.collect()\n", " \n", " # Process mask\n", " mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)\n", " mask = mask.transpose(1, 2, 0)\n", " mask_target = cv2.resize(mask, target_size)\n", " mask_target = (mask_target * 255).astype(np.uint8)\n", " \n", " # === FIX: Preserve original colors ===\n", " # Get the original image resized to target size\n", " if isinstance(original_image, np.ndarray):\n", " image_resized = cv2.resize(original_image, target_size)\n", " else:\n", " image_resized = np.array(original_image.resize(target_size, Image.Resampling.LANCZOS))\n", " \n", " # Create result by copying the resized original\n", " result = image_resized.copy()\n", " \n", " # Handle mask dimensions\n", " if len(mask_target.shape) == 3 and mask_target.shape[2] == 1:\n", " mask_target = mask_target.squeeze()\n", " \n", " # Set background to white\n", " result[mask_target < 128] = [255, 255, 255]\n", " \n", " return result, mask_target\n", " \n", " except Exception as e:\n", " print(f\"Warning: Background removal failed - {e}\")\n", " if isinstance(image, Image.Image):\n", " return np.array(image.resize(target_size)), None\n", " return cv2.resize(image, target_size), None\n", " \n", " def __del__(self):\n", " try:\n", " gc.collect()\n", " except:\n", " pass\n", "'''\n", "\n", "# Save to utils folder\n", "utils_dir = os.path.join(current_dir, 'utils')\n", "os.makedirs(utils_dir, exist_ok=True)\n", "\n", "with open(os.path.join(utils_dir, 'background_removal.py'), 'w', encoding='utf-8') as f:\n", " f.write(portable_code)\n", "print(f\"\\nโœ… Saved portable BackgroundRemover to {os.path.join(utils_dir, 'background_removal.py')}\")\n", "\n", "# %% [markdown]\n", "# ### Create setup script for deployment\n", "\n", "# %%\n", "setup_code = '''\"\"\"\n", "Setup script for U2Net\n", "Run this once before deploying to download U2Net\n", "\"\"\"\n", "import os\n", "import urllib.request\n", "import zipfile\n", "import shutil\n", "\n", "def setup_u2net():\n", " \"\"\"Download and setup U2Net\"\"\"\n", " \n", " # Create directories\n", " os.makedirs('u2net/model', exist_ok=True)\n", " os.makedirs('u2net/weights', exist_ok=True)\n", " \n", " # Download U2Net source\n", " print(\"๐Ÿ“ฅ Downloading U2Net source...\")\n", " u2net_url = \"https://github.com/xuebinqin/U-2-Net/archive/master.zip\"\n", " zip_path = \"u2net/u2net_source.zip\"\n", " \n", " try:\n", " urllib.request.urlretrieve(u2net_url, zip_path)\n", " \n", " # Extract\n", " with zipfile.ZipFile(zip_path, 'r') as zip_ref:\n", " zip_ref.extractall(\"u2net/\")\n", " \n", " # Copy model files\n", " source_model = os.path.join(\"u2net\", \"U-2-Net-master\", \"model\")\n", " if os.path.exists(source_model):\n", " for file in os.listdir(source_model):\n", " if file.endswith('.py'):\n", " shutil.copy(\n", " os.path.join(source_model, file),\n", " os.path.join(\"u2net\", \"model\", file)\n", " )\n", " \n", " # Clean up\n", " os.remove(zip_path)\n", " shutil.rmtree(os.path.join(\"u2net\", \"U-2-Net-master\"))\n", " \n", " print(\"\\nโœ… U2Net source setup complete!\")\n", " print(\"\\n๐Ÿ“Œ Next step:\")\n", " print(\"1. Download u2net.pth weights from:\")\n", " print(\" https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ\")\n", " print(\"2. Place it in: u2net/weights/u2net.pth\")\n", " \n", " except Exception as e:\n", " print(f\"โŒ Setup failed: {e}\")\n", "\n", "if __name__ == \"__main__\":\n", " setup_u2net()\n", "'''\n", "\n", "with open(os.path.join(current_dir, 'setup_u2net.py'), 'w', encoding='utf-8') as f:\n", " f.write(setup_code)\n", "print(\"โœ… Created setup_u2net.py - run this once before deploying\")\n", "\n", "# %% [markdown]\n", "# ### Create requirements.txt with all dependencies\n", "\n", "# %%\n", "requirements = '''torch>=1.9.0\n", "torchvision>=0.10.0\n", "numpy>=1.21.0\n", "opencv-python>=4.5.0\n", "Pillow>=8.0.0\n", "scikit-learn>=0.24.0\n", "matplotlib>=3.3.0\n", "flask>=2.0.0\n", "werkzeug>=2.0.0\n", "gdown # For downloading weights\n", "'''\n", "\n", "with open(os.path.join(current_dir, 'requirements.txt'), 'w', encoding='utf-8') as f:\n", " f.write(requirements)\n", "print(\"โœ… Created requirements.txt\")\n", "\n", "# %% [markdown]\n", "# ### Final Verification\n", "\n", "# %%\n", "print(\"\\n\" + \"=\"*70)\n", "print(\"๐ŸŽฏ PORTABLE BACKGROUND REMOVER - FINAL CHECK\")\n", "print(\"=\"*70)\n", "\n", "print(\"\\nโœ… Created files:\")\n", "print(\" โ€ข utils/background_removal.py - Portable U2Net class\")\n", "print(\" โ€ข setup_u2net.py - Run once to download U2Net source\")\n", "print(\" โ€ข requirements.txt - All dependencies\")\n", "\n", "print(\"\\n๐Ÿ“ Current folder structure:\")\n", "print(\" NitroSense-AI/\")\n", "print(\" โ”œโ”€โ”€ app.py\")\n", "print(\" โ”œโ”€โ”€ requirements.txt\")\n", "print(\" โ”œโ”€โ”€ setup_u2net.py\")\n", "print(\" โ”œโ”€โ”€ u2net/\")\n", "print(\" โ”‚ โ”œโ”€โ”€ model/\")\n", "print(\" โ”‚ โ”‚ โ””โ”€โ”€ u2net.py (source code)\")\n", "print(\" โ”‚ โ””โ”€โ”€ weights/\")\n", "print(f\" โ”‚ โ””โ”€โ”€ u2net.pth {'โœ…' if os.path.exists(weights_path) else 'โŒ'} (weights)\")\n", "print(\" โ””โ”€โ”€ utils/\")\n", "print(\" โ””โ”€โ”€ background_removal.py\")\n", "\n", "if os.path.exists(weights_path):\n", " print(\"\\nโœ… U2Net weights are present! Background removal will work.\")\n", " print(f\"\\n๐Ÿ“ Test results saved to: {os.path.join(current_dir, 'test_output_with_weights')}\")\n", " print(\" Check the following files:\")\n", " print(\" - 6538a45fe4f02_original.jpg (original resized image)\")\n", " print(\" - 6538a45fe4f02_mask.jpg (U2Net mask)\")\n", " print(\" - 6538a45fe4f02_bg_removed.jpg (final result)\")\n", "else:\n", " print(\"\\nโš ๏ธ U2Net weights are missing!\")\n", " print(\" Download from: https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ\")\n", " print(f\" Save to: {weights_path}\")\n", "\n", "print(\"\\n๐Ÿ“‹ Deployment steps:\")\n", "print(\"1. Run: python setup_u2net.py (downloads U2Net source)\")\n", "print(\"2. Ensure u2net.pth is in u2net/weights/\")\n", "print(\"3. Install: pip install -r requirements.txt\")\n", "print(\"4. Run: python app.py\")\n", "\n", "print(\"\\nโœ… Your app is now fully portable!\")\n", "print(\" Works on any system - no hardcoded paths!\")" ] }, { "cell_type": "code", "execution_count": null, "id": "403b4b1d", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.5" } }, "nbformat": 4, "nbformat_minor": 5 }