nermadie/2.5D_Depth_Studio

1"""2Script to test and compare quality across different settings3"""4 5import cv26import numpy as np7from PIL import Image8import matplotlib.pyplot as plt9import time10from pathlib import Path11 12 13class QualityComparator:14    def __init__(self):15        self.results = {}16 17    def compare_depth_methods(self, image_path):18        """Compare depth post-processing methods."""19 20        image = Image.open(image_path).convert("RGB")21        image_np = np.array(image)22 23        # Assume we already have a depth-like signal to compare filters24        gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)25 26        results = {}27 28        # Method 1: Simple Gaussian29        start = time.time()30        depth1 = cv2.GaussianBlur(gray, (5, 5), 0)31        results["gaussian"] = {"depth": depth1, "time": time.time() - start}32 33        # Method 2: Bilateral Filter34        start = time.time()35        depth2 = cv2.bilateralFilter(gray, 9, 75, 75)36        results["bilateral"] = {"depth": depth2, "time": time.time() - start}37 38        # Method 3: Edge Preserving39        start = time.time()40        depth3 = cv2.edgePreservingFilter(gray, flags=1, sigma_s=60, sigma_r=0.4)41        results["edge_preserving"] = {"depth": depth3, "time": time.time() - start}42 43        # Visualize44        fig, axes = plt.subplots(2, 2, figsize=(12, 10))45 46        axes[0, 0].imshow(image)47        axes[0, 0].set_title("Original Image")48        axes[0, 0].axis("off")49 50        for idx, (name, data) in enumerate(results.items(), 1):51            row = idx // 252            col = idx % 253            axes[row, col].imshow(data["depth"], cmap="magma")54            axes[row, col].set_title(f'{name.title()}\nTime: {data["time"]:.3f}s')55            axes[row, col].axis("off")56 57        plt.tight_layout()58        plt.savefig("depth_comparison.png", dpi=150, bbox_inches="tight")59        print("✅ Depth comparison saved to depth_comparison.png")60 61        return results62 63    def compare_inpainting(self, image_np, mask):64        """Compare inpainting methods."""65 66        results = {}67 68        # Method 1: Navier-Stokes69        start = time.time()70        result1 = cv2.inpaint(image_np, mask, 5, cv2.INPAINT_NS)71        results["navier_stokes"] = {"result": result1, "time": time.time() - start}72 73        # Method 2: TELEA74        start = time.time()75        result2 = cv2.inpaint(image_np, mask, 5, cv2.INPAINT_TELEA)76        results["telea"] = {"result": result2, "time": time.time() - start}77 78        # Method 3: Multi-scale TELEA79        start = time.time()80        result3 = self._multiscale_inpaint(image_np, mask)81        results["multiscale"] = {"result": result3, "time": time.time() - start}82 83        # Visualize84        fig, axes = plt.subplots(2, 2, figsize=(12, 10))85 86        # Show original with mask overlay87        masked_img = image_np.copy()88        masked_img[mask > 0] = [255, 0, 0]  # Red for masked area89        axes[0, 0].imshow(masked_img)90        axes[0, 0].set_title("Original + Mask")91        axes[0, 0].axis("off")92 93        for idx, (name, data) in enumerate(results.items(), 1):94            row = idx // 295            col = idx % 296            axes[row, col].imshow(data["result"])97            axes[row, col].set_title(98                f'{name.replace("_", " ").title()}\nTime: {data["time"]:.3f}s'99            )100            axes[row, col].axis("off")101 102        plt.tight_layout()103        plt.savefig("inpaint_comparison.png", dpi=150, bbox_inches="tight")104        print("✅ Inpaint comparison saved to inpaint_comparison.png")105 106        return results107 108    def _multiscale_inpaint(self, image, mask):109        """Helper: multi-scale inpainting"""110        scales = [1.0, 0.5, 0.25]111        results = []112 113        for scale in scales:114            h, w = int(image.shape[0] * scale), int(image.shape[1] * scale)115            img_scaled = cv2.resize(image, (w, h))116            mask_scaled = cv2.resize(mask, (w, h))117 118            inpainted = cv2.inpaint(img_scaled, mask_scaled, 5, cv2.INPAINT_TELEA)119            inpainted = cv2.resize(inpainted, (image.shape[1], image.shape[0]))120            results.append(inpainted)121 122        final = results[0] * 0.5 + results[1] * 0.3 + results[2] * 0.2123        return final.astype(np.uint8)124 125    def compare_soft_masks(self, hard_mask):126        """Compare different ways to generate soft masks."""127 128        results = {}129 130        # Method 1: Simple Gaussian131        start = time.time()132        mask1 = cv2.GaussianBlur(hard_mask.astype(np.float32), (21, 21), 0)133        results["gaussian"] = {"mask": mask1, "time": time.time() - start}134 135        # Method 2: Morphology + Gaussian136        start = time.time()137        kernel = np.ones((5, 5), np.uint8)138        mask2 = cv2.morphologyEx(hard_mask, cv2.MORPH_CLOSE, kernel, iterations=2)139        mask2 = cv2.GaussianBlur(mask2.astype(np.float32), (15, 15), 0)140        results["morph_gaussian"] = {"mask": mask2, "time": time.time() - start}141 142        # Method 3: Distance Transform143        start = time.time()144        dist = cv2.distanceTransform(hard_mask, cv2.DIST_L2, 5)145        mask3 = cv2.normalize(dist, None, 0, 1, cv2.NORM_MINMAX)146        mask3 = cv2.GaussianBlur(mask3, (11, 11), 0)147        results["distance_transform"] = {"mask": mask3, "time": time.time() - start}148 149        # Visualize150        fig, axes = plt.subplots(2, 2, figsize=(12, 10))151 152        axes[0, 0].imshow(hard_mask, cmap="gray")153        axes[0, 0].set_title("Hard Mask (Original)")154        axes[0, 0].axis("off")155 156        for idx, (name, data) in enumerate(results.items(), 1):157            row = idx // 2158            col = idx % 2159            axes[row, col].imshow(data["mask"], cmap="gray")160            axes[row, col].set_title(161                f'{name.replace("_", " ").title()}\nTime: {data["time"]:.3f}s'162            )163            axes[row, col].axis("off")164 165        plt.tight_layout()166        plt.savefig("mask_comparison.png", dpi=150, bbox_inches="tight")167        print("✅ Mask comparison saved to mask_comparison.png")168 169        return results170 171    def benchmark_full_pipeline(self, image_path, configs):172        """Test the full pipeline with different configs."""173 174        print("🚀 Starting benchmark...")175        results = {}176 177        for name, config in configs.items():178            print(f"\n📊 Testing: {name}")179            start = time.time()180 181            # Simulate processing182            # In a real scenario, call the actual processing functions183            time.sleep(1)  # Placeholder184 185            total_time = time.time() - start186            results[name] = {"time": total_time, "config": config}187 188            print(f"   Time: {total_time:.2f}s")189 190        # Print summary191        print("\n" + "=" * 50)192        print("BENCHMARK SUMMARY")193        print("=" * 50)194 195        for name, data in sorted(results.items(), key=lambda x: x[1]["time"]):196            print(f"{name:20s} {data['time']:8.2f}s")197 198        return results199 200    def quality_metrics(self, original, processed):201        """Compute metrics to evaluate quality."""202 203        # Convert to grayscale for metrics204        if len(original.shape) == 3:205            orig_gray = cv2.cvtColor(original, cv2.COLOR_RGB2GRAY)206            proc_gray = cv2.cvtColor(processed, cv2.COLOR_RGB2GRAY)207        else:208            orig_gray = original209            proc_gray = processed210 211        # 1. PSNR (Peak Signal-to-Noise Ratio)212        mse = np.mean((orig_gray - proc_gray) ** 2)213        if mse == 0:214            psnr = 100215        else:216            psnr = 20 * np.log10(255.0 / np.sqrt(mse))217 218        # 2. SSIM (Structural Similarity Index)219        from skimage.metrics import structural_similarity as ssim220 221        ssim_value = ssim(orig_gray, proc_gray)222 223        # 3. Edge preservation224        edges_orig = cv2.Canny(orig_gray, 50, 150)225        edges_proc = cv2.Canny(proc_gray, 50, 150)226        edge_similarity = np.sum(edges_orig == edges_proc) / edges_orig.size227 228        return {"PSNR": psnr, "SSIM": ssim_value, "Edge Preservation": edge_similarity}229 230 231# ============================================================================232# USAGE EXAMPLES233# ============================================================================234 235if __name__ == "__main__":236    comparator = QualityComparator()237 238    # Example 1: Compare depth processing239    print("📸 Test 1: Depth Processing Methods")240    print("-" * 50)241    # comparator.compare_depth_methods("test_image.jpg")242 243    # Example 2: Compare inpainting244    print("\n🎨 Test 2: Inpainting Methods")245    print("-" * 50)246    # Load sample image and mask247    # image = cv2.imread("test_image.jpg")248    # mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)249    # mask[100:200, 100:200] = 255  # Sample mask250    # comparator.compare_inpainting(image, mask)251 252    # Example 3: Compare soft masks253    print("\n✨ Test 3: Soft Mask Methods")254    print("-" * 50)255    # hard_mask = np.zeros((400, 400), dtype=np.uint8)256    # cv2.circle(hard_mask, (200, 200), 100, 255, -1)257    # comparator.compare_soft_masks(hard_mask)258 259    # Example 4: Full pipeline benchmark260    print("\n⚡ Test 4: Full Pipeline Benchmark")261    print("-" * 50)262 263    configs = {264        "Mobile": {"resize": 640, "model": "hybrid", "layers": 3},265        "Balanced": {"resize": 1024, "model": "large", "layers": 4},266        "Quality": {"resize": 1920, "model": "large", "layers": 4},267    }268 269    # comparator.benchmark_full_pipeline("test_image.jpg", configs)270 271    print("\n✅ All tests complete! Check output images.")272    print("\n💡 Tips:")273    print("   - Use bilateral filter for depth (best edge preservation)")274    print("   - Use TELEA for inpainting (better than NS)")275    print("   - Use morphology + gaussian for soft masks")276    print("   - Quality setting for best results, Balanced for speed")277