app.py added both detection

app.py

@@ -19,17 +19,32 @@ def load_model(repo_id):
 traffic_cones_model = load_model("ExStella/Traffic-cones")
 license_plate_model = load_model("ExStella/License-plate")
 
-# Function to process an image with the selected model
+# Function to process an image with the selected model(s)
 def process_image(img, model_type):
-    model = traffic_cones_model if model_type == "Traffic Cones" else license_plate_model
+    if model_type == "Traffic Cones":
+        model = traffic_cones_model
+    elif model_type == "License Plate":
+        model = license_plate_model
+    elif model_type == "Both":
+        # Process with both models
+        result1 = traffic_cones_model.predict(img, conf=0.5, iou=0.6)
+        result2 = license_plate_model.predict(img, conf=0.5, iou=0.6)
+        img_bgr1 = result1[0].plot()
+        img_bgr2 = result2[0].plot()
+        img_bgr_combined = cv2.addWeighted(img_bgr1, 0.5, img_bgr2, 0.5, 0)  # Combine results
+        output_image = Image.fromarray(img_bgr_combined[..., ::-1])  # Convert BGR to RGB
+        return output_image
+    else:
+        raise ValueError("Invalid detection type.")
+    
+    # Process for single model type
     result = model.predict(img, conf=0.5, iou=0.6)
     img_bgr = result[0].plot()
-    output_image = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB for PIL
+    output_image = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB
     return output_image
 
-# Function to process a video with the selected model
+# Function to process a video with the selected model(s)
 def process_video(video_path, model_type):
-    model = traffic_cones_model if model_type == "Traffic Cones" else license_plate_model
     temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
     cap = cv2.VideoCapture(video_path)
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -45,9 +60,24 @@ def process_video(video_path, model_type):
         if not ret:
             break
         
-        # Perform object detection
-        results = model.predict(frame, conf=0.5, iou=0.6)
-        frame = results[0].plot()  # Annotated frame
+        if model_type == "Traffic Cones":
+            model = traffic_cones_model
+        elif model_type == "License Plate":
+            model = license_plate_model
+        elif model_type == "Both":
+            # Process with both models
+            results1 = traffic_cones_model.predict(frame, conf=0.5, iou=0.6)
+            results2 = license_plate_model.predict(frame, conf=0.5, iou=0.6)
+            frame1 = results1[0].plot()
+            frame2 = results2[0].plot()
+            frame = cv2.addWeighted(frame1, 0.5, frame2, 0.5, 0)  # Combine results
+        else:
+            raise ValueError("Invalid detection type.")
+        
+        # Annotate frame for single model type
+        if model_type in ["Traffic Cones", "License Plate"]:
+            results = model.predict(frame, conf=0.5, iou=0.6)
+            frame = results[0].plot()  # Annotated frame
         
         # Write the processed frame
         out.write(frame)
@@ -59,6 +89,9 @@ def process_video(video_path, model_type):
 
 # Gradio prediction function
 def predict(input_file, model_type):
+    if not model_type:
+        raise ValueError("Please select a detection type before submitting.")
+    
     if input_file.name.endswith(('.jpg', '.jpeg', '.png')):
         # Image input
         img = Image.open(input_file.name)
@@ -76,12 +109,12 @@ gr.Interface(
     fn=predict,
     inputs=[
         gr.File(label="Upload an image or video (JPG, PNG, MP4, AVI, etc.)"),
-        gr.Radio(["Traffic Cones", "License Plate"], label="Choose Detection Type"),
+        gr.Radio(["Traffic Cones", "License Plate", "Both"], label="Choose Detection Type"),
     ],
     outputs=[
         gr.Image(type="pil", label="Processed Image"),  # Output for images
         gr.Video(label="Processed Video"),  # Output for videos
     ],
     title="Object Detection for Traffic Cones and License Plates",
-    description="Upload an image or video to perform object detection. Select between Traffic Cones or License Plates detection.",
-).launch(share=True)
+    description="Upload an image or video to perform object detection. Select between Traffic Cones, License Plates, or Both detection types.",
+).launch(share=True)