Upload app.py

app.py

@@ -31,37 +31,52 @@ metrics_df = pd.DataFrame(metrics_data, columns=["Metric", "Value"])
 ##################################
 
 def run_inference(img: np.ndarray, model):
+    """
+    Runs inference on the input image using the YOLO model.
+    Returns the annotated image, Top-1 prediction, and Top-5 predictions.
+    """
+    # Convert from BGR to RGB
     img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    results = model.predict(img_rgb)  # Run prediction
+    
+    # Run prediction
+    results = model.predict(img_rgb)
+    
+    # Extract probabilities (if available)
     result_probs = results[0].probs
-
+    
     # Get top-1 and top-5 predictions
     top1_class = result_probs.top1
     top5_classes = result_probs.top5
     top1_conf = result_probs.top1conf.item()
     top5_conf = result_probs.top5conf
-
-    # Generate annotated image
-    annotated_img = results[0].plot()
-
+    
+    # Generate annotated image (RGB format)
+    annotated_img = results[0].plot()  # Assuming this returns RGB
+    
     # Format results
     top1_result = f"Class: {model.names[top1_class]}, Confidence: {top1_conf:.2f}"
     top5_results = [
         f"{model.names[c]}: {conf:.2f}" for c, conf in zip(top5_classes, top5_conf)
     ]
-
+    
     return annotated_img, top1_result, top5_results
 
 def process_image(image):
+    """
+    Processes the input image, runs inference, and prepares the outputs.
+    """
+    # Convert PIL Image to NumPy array
     img = np.array(image)
+    
+    # Convert from RGB to BGR for OpenCV (if needed by YOLO)
     img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-
+    
     # Run classification inference
     annotated_img, top1_result, top5_results = run_inference(img_bgr, model)
-
-    # Convert annotated image back to PIL format
-    annotated_img_pil = Image.fromarray(cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB))
-
+    
+    # Convert annotated image back to PIL format without altering color channels
+    annotated_img_pil = Image.fromarray(annotated_img)  # Assuming annotated_img is in RGB
+    
     # Return the annotated image, Top-1, and Top-5 predictions, along with metrics
     return annotated_img_pil, f"Top-1: {top1_result}", "\n".join(top5_results), metrics_df
 
@@ -74,7 +89,7 @@ with gr.Blocks() as demo:
             input_image = gr.Image(type="pil", label="Upload Image")
             submit_btn = gr.Button("Run Inference")
         with gr.Column():
-            annotated_image = gr.Image(type="pil", label="Annotated Image")  # Updated to show annotated image
+            annotated_image = gr.Image(type="pil", label="Annotated Image")  # Shows annotated image in RGB
             top1_output = gr.Textbox(label="Top-1 Prediction")
             top5_output = gr.Textbox(label="Top-5 Predictions")
             metrics_table = gr.DataFrame(value=metrics_df, label="Validation Metrics")