Update app.py

app.py

@@ -7,27 +7,32 @@ def predict(img):
 
     model = YOLO("model.pt")
 
-    results = model.predict(source=img,save=False, show_labels=False, show_conf=False)
+    results = model.predict(source=img, save=False, show_labels=False, show_conf=False)
+    
+    # Handle both single result and list of results
+    if not isinstance(results, (list, tuple)):
+        results = [results]
+    
     count = 0
+    annot_img = None
     for i in results:
         count = len(i.boxes)
         annot_img = i.plot(labels=False, masks=False)
+        # Convert BGR to RGB if needed
+        if annot_img.shape[-1] == 3:  # Check if image has 3 channels
+            annot_img = annot_img[..., ::-1]  # Reverse the color channels
     
-    #result_img = os.path.join(cur_path,r"runs\segment\predict",f"{path}.jpg")
-    #print("result img",result_img)
     return str(count), annot_img
 
 with gr.Blocks(title="Pill Counter") as demo:
-
     with gr.Row():
         with gr.Column():
-            img = gr.Image(type="filepath",format="jpg", height=500, width=700)
-
+            img = gr.Image(type="filepath", format="jpg", height=500, width=700)
             button = gr.Button()
         with gr.Column():
             data_output = gr.Textbox()
             img_output = gr.Image(type="numpy")
-        button.click(fn=predict, inputs=img, outputs=[data_output,img_output])
+        button.click(fn=predict, inputs=img, outputs=[data_output, img_output])
 
 if __name__ == "__main__":
     demo.launch()