Update app.py

app.py

@@ -12,8 +12,8 @@ def load_model():
 
 ort_session = load_model()
 
-# Define class names (update this based on your model's classes)
-CLASS_NAMES = ['car', 'license_plate']
+# Define class names and their corresponding indices
+CLASS_NAMES = {0: 'car', 1: 'license_plate'}
 
 def preprocess_image(image, target_size=(640, 640)):
     if isinstance(image, Image.Image):
@@ -57,9 +57,9 @@ def postprocess_results(output, image_shape, confidence_threshold=0.25, iou_thre
     for i in indices:
         box = boxes[i]
         score = scores[i]
-        class_id = class_ids[i]
+        class_id = int(class_ids[i])
         x1, y1, x2, y2 = map(int, box)
-        results.append((x1, y1, x2, y2, float(score), int(class_id)))
+        results.append((x1, y1, x2, y2, float(score), class_id))
 
     return results
 
@@ -73,9 +73,9 @@ def process_image(image):
     results = postprocess_results(outputs, image.shape)
     
     for x1, y1, x2, y2, score, class_id in results:
-        color = (0, 255, 0) if CLASS_NAMES[class_id] == 'car' else (255, 0, 0)
+        color = (0, 255, 0) if CLASS_NAMES.get(class_id, 'unknown') == 'car' else (255, 0, 0)
         cv2.rectangle(orig_image, (x1, y1), (x2, y2), color, 2)
-        label = f"{CLASS_NAMES[class_id]}: {score:.2f}"
+        label = f"{CLASS_NAMES.get(class_id, 'unknown')}: {score:.2f}"
         cv2.putText(orig_image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
     
     return cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)