Update app.py

app.py

@@ -113,99 +113,89 @@ class ImageAnalyzer:
                     "size": image.size,
                     "aspect_ratio": image.size[0] / image.size[1]
                 }
-
-        # Edge detection for crack analysis
-        gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
-        edges = cv2.Canny(gray, 100, 200)
-        stats["edge_density"] = np.mean(edges > 0)
-
-        # Color analysis for rust detection
-        hsv = cv2.cvtColor(img_array, cv2.COLOR_RGB2HSV)
-        rust_mask = cv2.inRange(hsv, np.array([0, 50, 50]), np.array([30, 255, 255]))
-        stats["rust_percentage"] = np.mean(rust_mask > 0)
-
-        # Transform for model
-        model_input = self.transforms(image).unsqueeze(0).to(self.device)
-
-        return {
-            "model_input": model_input,
-            "stats": stats,
-            "edges": edges,
-            "rust_mask": rust_mask
-        }
-except Exception as e:
-        logger.error(f"Preprocessing error: {e}")
-        return None
-
-
-    def detect_defects(self, image: Image.Image) -> Dict[str, Any]:
-    """Enhanced defect detection with multiple analysis methods"""
-    try:
-        # Preprocess image
-        proc_data = self.preprocess_image(image)
-        if proc_data is None:
-            logger.error("Image preprocessing failed.")
-            return None  # Early return if preprocessing failed
-
-        # Model prediction
-        with torch.no_grad():
-            outputs = self.model(proc_data["model_input"])
+                # Edge detection for crack analysis
+                gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
+                edges = cv2.Canny(gray, 100, 200)
+                stats["edge_density"] = np.mean(edges > 0)
+                # Color analysis for rust detection
+                hsv = cv2.cvtColor(img_array, cv2.COLOR_RGB2HSV)
+                rust_mask = cv2.inRange(hsv, np.array([0, 50, 50]), np.array([30, 255, 255]))
+                stats["rust_percentage"] = np.mean(rust_mask > 0)
+                # Transform for model
+                model_input = self.transforms(image).unsqueeze(0).to(self.device)
+                return {
+                    "model_input": model_input,
+                    "stats": stats,
+                    "edges": edges,
+                    "rust_mask": rust_mask
+                }
+        except Exception as e:
+            logger.error(f"Preprocessing error: {e}")
+            return None
+            def detect_defects(self, image: Image.Image) -> Dict[str, Any]:
+                """Enhanced defect detection with multiple analysis methods"""
+                try:
+                    # Preprocess image
+                    proc_data = self.preprocess_image(image)
+                    if proc_data is None:
+                        logger.error("Image preprocessing failed.")
+                        return None  # Early return if preprocessing failed
             
-        # Get probabilities
-        probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
-        
-        # Convert to dictionary
-        defect_probs = {
-            self.defect_classes[i]: float(probabilities[0][i])
-            for i in range(len(self.defect_classes))
-        }
-
-        # Generate attention heatmap
-        attention_weights = outputs.attentions[-1].mean(dim=1)[0] if hasattr(outputs, 'attentions') else None
-        heatmap = self.generate_heatmap(attention_weights, image.size) if attention_weights is not None else None
-
-        # Additional analysis based on image statistics
-        additional_analysis = self.analyze_image_statistics(proc_data["stats"])
-
-        # Combine all results
-        result = {
-            "defect_probabilities": defect_probs,
-            "heatmap": heatmap,
-            "image_statistics": proc_data["stats"],
-            "additional_analysis": additional_analysis,
-            "edge_detection": proc_data["edges"],
-            "rust_detection": proc_data["rust_mask"],
-            "timestamp": datetime.now().isoformat()
-        }
-
-        # Save to history
-        self.history.append(result)
-
-        return result
-    except Exception as e:
-        logger.error(f"Defect detection error: {e}")
-        return None
-
-
-    def analyze_image_statistics(self, stats: Dict) -> Dict[str, Any]:
-        """Analyze image statistics for additional insights"""
-        analysis = {}
-        
-        # Brightness analysis
-        if stats["mean_brightness"] < 50:
-            analysis["lighting_condition"] = "Poor lighting - may affect accuracy"
-        elif stats["mean_brightness"] > 200:
-            analysis["lighting_condition"] = "Overexposed - may affect accuracy"
-        
-        # Edge density analysis
-        if stats["edge_density"] > 0.1:
-            analysis["crack_likelihood"] = "High crack probability based on edge detection"
-        
-        # Rust analysis
-        if stats["rust_percentage"] > 0.05:
-            analysis["corrosion_indicator"] = "Possible corrosion detected"
-
-        return analysis
+                    # Model prediction
+                    with torch.no_grad():
+                        outputs = self.model(proc_data["model_input"])
+                        
+                    # Get probabilities
+                    probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
+                    
+                    # Convert to dictionary
+                    defect_probs = {
+                        self.defect_classes[i]: float(probabilities[0][i])
+                        for i in range(len(self.defect_classes))
+                    }
+                    # Generate attention heatmap
+                    attention_weights = outputs.attentions[-1].mean(dim=1)[0] if hasattr(outputs, 'attentions') else None
+                    heatmap = self.generate_heatmap(attention_weights, image.size) if attention_weights is not None else None
+            
+                    # Additional analysis based on image statistics
+                    additional_analysis = self.analyze_image_statistics(proc_data["stats"])
+            
+                    # Combine all results
+                    result = {
+                        "defect_probabilities": defect_probs,
+                        "heatmap": heatmap,
+                        "image_statistics": proc_data["stats"],
+                        "additional_analysis": additional_analysis,
+                        "edge_detection": proc_data["edges"],
+                        "rust_detection": proc_data["rust_mask"],
+                        "timestamp": datetime.now().isoformat()
+                    }
+                    # Save to history
+                    self.history.append(result)
+                    return result
+                except Exception as e:
+                    logger.error(f"Defect detection error: {e}")
+                    return None
+                    
+                    def analyze_image_statistics(self, stats: Dict) -> Dict[str, Any]:
+                        """Analyze image statistics for additional insights"""
+                        analysis = {}
+                        
+                        # Brightness analysis
+                        if stats["mean_brightness"] < 50:
+                            analysis["lighting_condition"] = "Poor lighting - may affect accuracy"
+                        elif stats["mean_brightness"] > 200:
+                            analysis["lighting_condition"] = "Overexposed - may affect accuracy"
+                        
+                        # Edge density analysis
+                        if stats["edge_density"] > 0.1:
+                            analysis["crack_likelihood"] = "High crack probability based on edge detection"
+                        
+                        # Rust analysis
+                        if stats["rust_percentage"] > 0.05:
+                            analysis["corrosion_indicator"] = "Possible corrosion detected"
+                
+                        return analysis
 
     def generate_heatmap(self, attention_weights: torch.Tensor, image_size: Tuple[int, int]) -> np.ndarray:
         """Generate enhanced attention heatmap"""