Víctor Sáez commited on
Commit
82026bf
·
1 Parent(s): 862343e

Add exception handling

Browse files
Files changed (1) hide show
  1. app.py +83 -86
app.py CHANGED
@@ -180,93 +180,90 @@ def translate_label(language_label, label):
180
 
181
  def detect_objects(image, language_selector, translated_model_selector, threshold):
182
  """Enhanced object detection with adjustable threshold and better info"""
183
- # Get the actual model key from the translated name
184
-
185
- def detect_objects(image, language_selector, translated_model_selector, threshold):
186
- try:
187
- if image is None:
188
- return None, "Por favor, sube una imagen antes de detectar objetos."
189
-
190
- model_selector = get_model_key_from_translation(translated_model_selector, language_selector)
191
- print(f"Processing image. Language: {language_selector}, Model: {model_selector}, Threshold: {threshold}")
192
-
193
- model, processor = load_model(model_selector)
194
-
195
- inputs = processor(images=image, return_tensors="pt")
196
- outputs = model(**inputs)
197
-
198
- target_sizes = torch.tensor([image.size[::-1]])
199
- results = processor.post_process_object_detection(
200
- outputs, threshold=threshold, target_sizes=target_sizes
201
- )[0]
202
-
203
- image_with_boxes = image.copy()
204
- draw = ImageDraw.Draw(image_with_boxes)
205
-
206
- detection_info = f"Detected {len(results['scores'])} objects with threshold {threshold}\n"
207
- detection_info += f"Model: {translated_model_selector} ({model_selector})\n\n"
208
-
209
- colors = {
210
- 'high': 'red', # > 0.8
211
- 'medium': 'orange', # 0.5-0.8
212
- 'low': 'yellow' # < 0.5
213
- }
214
-
215
- detected_objects = []
216
-
217
- for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
218
- confidence = score.item()
219
- box = [round(x, 2) for x in box.tolist()]
220
- if confidence > 0.8:
221
- color = colors['high']
222
- elif confidence > 0.5:
223
- color = colors['medium']
224
- else:
225
- color = colors['low']
226
-
227
- draw.rectangle(box, outline=color, width=3)
228
- label_text = model.config.id2label[label.item()]
229
- translated_label = translate_label(language_selector, label_text)
230
- display_text = f"{translated_label}: {round(confidence, 3)}"
231
- detected_objects.append({
232
- 'label': label_text,
233
- 'translated': translated_label,
234
- 'confidence': confidence,
235
- 'box': box
236
- })
237
-
238
- try:
239
- image_width = image.size[0]
240
- font_size = max(image_width // 40, 12)
241
- font = get_font(font_size)
242
- text_bbox = draw.textbbox((0, 0), display_text, font=font)
243
- text_width = text_bbox[2] - text_bbox[0]
244
- text_height = text_bbox[3] - text_bbox[1]
245
- except:
246
- font = get_font(12)
247
- text_width = 50
248
- text_height = 20
249
-
250
- text_bg = [
251
- box[0], box[1] - text_height - 4,
252
- box[0] + text_width + 4, box[1]
253
- ]
254
- draw.rectangle(text_bg, fill="black")
255
- draw.text((box[0] + 2, box[1] - text_height - 2), display_text, fill="white", font=font)
256
-
257
- if detected_objects:
258
- detection_info += "Objects found:\n"
259
- for obj in sorted(detected_objects, key=lambda x: x['confidence'], reverse=True):
260
- detection_info += f"- {obj['translated']} ({obj['label']}): {obj['confidence']:.3f}\n"
261
  else:
262
- detection_info += "No objects detected. Try lowering the threshold."
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
263
 
264
- return image_with_boxes, detection_info
265
- except Exception as e:
266
- import traceback
267
- print("ERROR EN DETECT_OBJECTS:", e)
268
- traceback.print_exc()
269
- return None, f"Error detecting objects: {e}"
 
 
 
 
 
 
 
270
 
271
 
272
  def build_app():
@@ -353,4 +350,4 @@ load_model("DETR ResNet-50")
353
  # Launch the application
354
  if __name__ == "__main__":
355
  app = build_app()
356
- app.launch()
 
180
 
181
  def detect_objects(image, language_selector, translated_model_selector, threshold):
182
  """Enhanced object detection with adjustable threshold and better info"""
183
+ try:
184
+ if image is None:
185
+ return None, "Por favor, sube una imagen antes de detectar objetos."
186
+
187
+ model_selector = get_model_key_from_translation(translated_model_selector, language_selector)
188
+ print(f"Processing image. Language: {language_selector}, Model: {model_selector}, Threshold: {threshold}")
189
+
190
+ model, processor = load_model(model_selector)
191
+
192
+ inputs = processor(images=image, return_tensors="pt")
193
+ outputs = model(**inputs)
194
+
195
+ target_sizes = torch.tensor([image.size[::-1]])
196
+ results = processor.post_process_object_detection(
197
+ outputs, threshold=threshold, target_sizes=target_sizes
198
+ )[0]
199
+
200
+ image_with_boxes = image.copy()
201
+ draw = ImageDraw.Draw(image_with_boxes)
202
+
203
+ detection_info = f"Detected {len(results['scores'])} objects with threshold {threshold}\n"
204
+ detection_info += f"Model: {translated_model_selector} ({model_selector})\n\n"
205
+
206
+ colors = {
207
+ 'high': 'red', # > 0.8
208
+ 'medium': 'orange', # 0.5-0.8
209
+ 'low': 'yellow' # < 0.5
210
+ }
211
+
212
+ detected_objects = []
213
+
214
+ for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
215
+ confidence = score.item()
216
+ box = [round(x, 2) for x in box.tolist()]
217
+ if confidence > 0.8:
218
+ color = colors['high']
219
+ elif confidence > 0.5:
220
+ color = colors['medium']
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
221
  else:
222
+ color = colors['low']
223
+
224
+ draw.rectangle(box, outline=color, width=3)
225
+ label_text = model.config.id2label[label.item()]
226
+ translated_label = translate_label(language_selector, label_text)
227
+ display_text = f"{translated_label}: {round(confidence, 3)}"
228
+ detected_objects.append({
229
+ 'label': label_text,
230
+ 'translated': translated_label,
231
+ 'confidence': confidence,
232
+ 'box': box
233
+ })
234
+
235
+ try:
236
+ image_width = image.size[0]
237
+ font_size = max(image_width // 40, 12)
238
+ font = get_font(font_size)
239
+ text_bbox = draw.textbbox((0, 0), display_text, font=font)
240
+ text_width = text_bbox[2] - text_bbox[0]
241
+ text_height = text_bbox[3] - text_bbox[1]
242
+ except:
243
+ font = get_font(12)
244
+ text_width = 50
245
+ text_height = 20
246
+
247
+ text_bg = [
248
+ box[0], box[1] - text_height - 4,
249
+ box[0] + text_width + 4, box[1]
250
+ ]
251
+ draw.rectangle(text_bg, fill="black")
252
+ draw.text((box[0] + 2, box[1] - text_height - 2), display_text, fill="white", font=font)
253
 
254
+ if detected_objects:
255
+ detection_info += "Objects found:\n"
256
+ for obj in sorted(detected_objects, key=lambda x: x['confidence'], reverse=True):
257
+ detection_info += f"- {obj['translated']} ({obj['label']}): {obj['confidence']:.3f}\n"
258
+ else:
259
+ detection_info += "No objects detected. Try lowering the threshold."
260
+
261
+ return image_with_boxes, detection_info
262
+ except Exception as e:
263
+ import traceback
264
+ print("ERROR EN DETECT_OBJECTS:", e)
265
+ traceback.print_exc()
266
+ return None, f"Error detecting objects: {e}"
267
 
268
 
269
  def build_app():
 
350
  # Launch the application
351
  if __name__ == "__main__":
352
  app = build_app()
353
+ app.launch()