Update app.py

app.py

@@ -21,6 +21,19 @@ COLORS = [
     [0.301, 0.745, 0.933]
 ]
 
+import numpy as np
+import tensorflow as tf
+
+# Load EV plate classifier
+ev_model = tf.keras.models.load_model("plate_color_model.h5")
+
+def is_green_plate(plate_img):
+    plate_img = plate_img.resize((128,128))
+    plate_img = np.array(plate_img)/255.0
+    plate_img = np.expand_dims(plate_img, axis=0)
+    pred = ev_model.predict(plate_img)[0][0]
+    return pred > 0.5
+
 def make_prediction(img, feature_extractor, model):
     inputs = feature_extractor(img, return_tensors="pt")
     outputs = model(**inputs)
@@ -44,23 +57,21 @@ def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
     boxes = output_dict["boxes"].tolist()
     scores = output_dict["scores"].tolist()
     labels = output_dict["labels"].tolist()
-    
+
     if id2label is not None:
-        
         labels = [id2label[x] for x in labels]
-        
 
-    plt.figure(figsize=(50, 50))
+    plt.figure(figsize=(20, 20))
     plt.imshow(img)
     ax = plt.gca()
     colors = COLORS * 100
+
     for score, (xmin, ymin, xmax, ymax), label, color in zip(scores, boxes, labels, colors):
-        if label == 'license-plates':
+        if score < threshold:
+            continue
 
-            # Crop plate
+        if label in ["license-plates", "Rego Plates"]:
             plate_crop = img.crop((xmin, ymin, xmax, ymax))
-
-            # Check EV
             ev = is_green_plate(plate_crop)
 
             if ev:
@@ -69,19 +80,23 @@ def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
             else:
                 plate_type = "Non-EV Plate"
                 box_color = "red"
+        else:
+            plate_type = label
+            box_color = "blue"
 
         ax.add_patch(
             plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
-                          fill=False, color=box_color, linewidth=10)
+                          fill=False, color=box_color, linewidth=3)
         )
 
         ax.text(
-            xmin, ymin - 20,
-            f"{plate_type} | {score:0.2f}",
-            fontsize=50,
+            xmin, ymin - 10,
+            f"{plate_type} | {score:.2f}",
+            fontsize=14,
             bbox=dict(facecolor=box_color, alpha=0.7),
             color="white"
         )
+
     plt.axis("off")
     return fig2img(plt.gcf())
     
@@ -111,7 +126,7 @@ def detect_objects(model_name,url_input,image_input,webcam_input,threshold):
         image = webcam_input
     
     #Make prediction
-    processed_outputs = make_prediction(image, feature_extractor, model)
+    processed_outputs = make_prediction(image.convert("RGB"), feature_extractor, model)
     
     #Visualize prediction
     viz_img = visualize_prediction(image, processed_outputs, threshold, model.config.id2label)