Shakir60 commited on
Commit
f26cc53
·
verified ·
1 Parent(s): e7c0374

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +82 -92
app.py CHANGED
@@ -113,99 +113,89 @@ class ImageAnalyzer:
113
  "size": image.size,
114
  "aspect_ratio": image.size[0] / image.size[1]
115
  }
116
-
117
- # Edge detection for crack analysis
118
- gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
119
- edges = cv2.Canny(gray, 100, 200)
120
- stats["edge_density"] = np.mean(edges > 0)
121
-
122
- # Color analysis for rust detection
123
- hsv = cv2.cvtColor(img_array, cv2.COLOR_RGB2HSV)
124
- rust_mask = cv2.inRange(hsv, np.array([0, 50, 50]), np.array([30, 255, 255]))
125
- stats["rust_percentage"] = np.mean(rust_mask > 0)
126
-
127
- # Transform for model
128
- model_input = self.transforms(image).unsqueeze(0).to(self.device)
129
-
130
- return {
131
- "model_input": model_input,
132
- "stats": stats,
133
- "edges": edges,
134
- "rust_mask": rust_mask
135
- }
136
- except Exception as e:
137
- logger.error(f"Preprocessing error: {e}")
138
- return None
139
-
140
-
141
- def detect_defects(self, image: Image.Image) -> Dict[str, Any]:
142
- """Enhanced defect detection with multiple analysis methods"""
143
- try:
144
- # Preprocess image
145
- proc_data = self.preprocess_image(image)
146
- if proc_data is None:
147
- logger.error("Image preprocessing failed.")
148
- return None # Early return if preprocessing failed
149
-
150
- # Model prediction
151
- with torch.no_grad():
152
- outputs = self.model(proc_data["model_input"])
153
 
154
- # Get probabilities
155
- probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
156
-
157
- # Convert to dictionary
158
- defect_probs = {
159
- self.defect_classes[i]: float(probabilities[0][i])
160
- for i in range(len(self.defect_classes))
161
- }
162
-
163
- # Generate attention heatmap
164
- attention_weights = outputs.attentions[-1].mean(dim=1)[0] if hasattr(outputs, 'attentions') else None
165
- heatmap = self.generate_heatmap(attention_weights, image.size) if attention_weights is not None else None
166
-
167
- # Additional analysis based on image statistics
168
- additional_analysis = self.analyze_image_statistics(proc_data["stats"])
169
-
170
- # Combine all results
171
- result = {
172
- "defect_probabilities": defect_probs,
173
- "heatmap": heatmap,
174
- "image_statistics": proc_data["stats"],
175
- "additional_analysis": additional_analysis,
176
- "edge_detection": proc_data["edges"],
177
- "rust_detection": proc_data["rust_mask"],
178
- "timestamp": datetime.now().isoformat()
179
- }
180
-
181
- # Save to history
182
- self.history.append(result)
183
-
184
- return result
185
- except Exception as e:
186
- logger.error(f"Defect detection error: {e}")
187
- return None
188
-
189
-
190
- def analyze_image_statistics(self, stats: Dict) -> Dict[str, Any]:
191
- """Analyze image statistics for additional insights"""
192
- analysis = {}
193
-
194
- # Brightness analysis
195
- if stats["mean_brightness"] < 50:
196
- analysis["lighting_condition"] = "Poor lighting - may affect accuracy"
197
- elif stats["mean_brightness"] > 200:
198
- analysis["lighting_condition"] = "Overexposed - may affect accuracy"
199
-
200
- # Edge density analysis
201
- if stats["edge_density"] > 0.1:
202
- analysis["crack_likelihood"] = "High crack probability based on edge detection"
203
-
204
- # Rust analysis
205
- if stats["rust_percentage"] > 0.05:
206
- analysis["corrosion_indicator"] = "Possible corrosion detected"
207
-
208
- return analysis
209
 
210
  def generate_heatmap(self, attention_weights: torch.Tensor, image_size: Tuple[int, int]) -> np.ndarray:
211
  """Generate enhanced attention heatmap"""
 
113
  "size": image.size,
114
  "aspect_ratio": image.size[0] / image.size[1]
115
  }
116
+ # Edge detection for crack analysis
117
+ gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
118
+ edges = cv2.Canny(gray, 100, 200)
119
+ stats["edge_density"] = np.mean(edges > 0)
120
+ # Color analysis for rust detection
121
+ hsv = cv2.cvtColor(img_array, cv2.COLOR_RGB2HSV)
122
+ rust_mask = cv2.inRange(hsv, np.array([0, 50, 50]), np.array([30, 255, 255]))
123
+ stats["rust_percentage"] = np.mean(rust_mask > 0)
124
+ # Transform for model
125
+ model_input = self.transforms(image).unsqueeze(0).to(self.device)
126
+ return {
127
+ "model_input": model_input,
128
+ "stats": stats,
129
+ "edges": edges,
130
+ "rust_mask": rust_mask
131
+ }
132
+ except Exception as e:
133
+ logger.error(f"Preprocessing error: {e}")
134
+ return None
135
+ def detect_defects(self, image: Image.Image) -> Dict[str, Any]:
136
+ """Enhanced defect detection with multiple analysis methods"""
137
+ try:
138
+ # Preprocess image
139
+ proc_data = self.preprocess_image(image)
140
+ if proc_data is None:
141
+ logger.error("Image preprocessing failed.")
142
+ return None # Early return if preprocessing failed
 
 
 
 
 
 
 
 
 
 
143
 
144
+ # Model prediction
145
+ with torch.no_grad():
146
+ outputs = self.model(proc_data["model_input"])
147
+
148
+ # Get probabilities
149
+ probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
150
+
151
+ # Convert to dictionary
152
+ defect_probs = {
153
+ self.defect_classes[i]: float(probabilities[0][i])
154
+ for i in range(len(self.defect_classes))
155
+ }
156
+ # Generate attention heatmap
157
+ attention_weights = outputs.attentions[-1].mean(dim=1)[0] if hasattr(outputs, 'attentions') else None
158
+ heatmap = self.generate_heatmap(attention_weights, image.size) if attention_weights is not None else None
159
+
160
+ # Additional analysis based on image statistics
161
+ additional_analysis = self.analyze_image_statistics(proc_data["stats"])
162
+
163
+ # Combine all results
164
+ result = {
165
+ "defect_probabilities": defect_probs,
166
+ "heatmap": heatmap,
167
+ "image_statistics": proc_data["stats"],
168
+ "additional_analysis": additional_analysis,
169
+ "edge_detection": proc_data["edges"],
170
+ "rust_detection": proc_data["rust_mask"],
171
+ "timestamp": datetime.now().isoformat()
172
+ }
173
+ # Save to history
174
+ self.history.append(result)
175
+ return result
176
+ except Exception as e:
177
+ logger.error(f"Defect detection error: {e}")
178
+ return None
179
+
180
+ def analyze_image_statistics(self, stats: Dict) -> Dict[str, Any]:
181
+ """Analyze image statistics for additional insights"""
182
+ analysis = {}
183
+
184
+ # Brightness analysis
185
+ if stats["mean_brightness"] < 50:
186
+ analysis["lighting_condition"] = "Poor lighting - may affect accuracy"
187
+ elif stats["mean_brightness"] > 200:
188
+ analysis["lighting_condition"] = "Overexposed - may affect accuracy"
189
+
190
+ # Edge density analysis
191
+ if stats["edge_density"] > 0.1:
192
+ analysis["crack_likelihood"] = "High crack probability based on edge detection"
193
+
194
+ # Rust analysis
195
+ if stats["rust_percentage"] > 0.05:
196
+ analysis["corrosion_indicator"] = "Possible corrosion detected"
197
+
198
+ return analysis
199
 
200
  def generate_heatmap(self, attention_weights: torch.Tensor, image_size: Tuple[int, int]) -> np.ndarray:
201
  """Generate enhanced attention heatmap"""