Update app.py

app.py

@@ -100,90 +100,93 @@ class ImageAnalyzer:
         ])
 
     def preprocess_image(self, image: Image.Image) -> Dict[str, Any]:
-        """Enhanced image preprocessing with multiple analyses"""
-        try:
-            # Convert to RGB if necessary
-            if image.mode != 'RGB':
-                image = image.convert('RGB')
-
-            # Basic image statistics
-            img_array = np.array(image)
-            stats = {
-                "mean_brightness": np.mean(img_array),
-                "std_brightness": np.std(img_array),
-                "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
+    """Enhanced image preprocessing with multiple analyses"""
+    try:
+        # Convert to RGB if necessary
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+
+        # Basic image statistics
+        img_array = np.array(image)
+        stats = {
+            "mean_brightness": np.mean(img_array),
+            "std_brightness": np.std(img_array),
+            "size": image.size,
+            "aspect_ratio": image.size[0] / image.size[1]
+        }
 
-    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:
-                return None
+        # 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)
 
-            # Model prediction
-            with torch.no_grad():
-                outputs = self.model(proc_data["model_input"])
-                
-            # Get probabilities
-            probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
+        # 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"])
             
-            # 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
+        # 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"""