Update sample.py

sample.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from PIL import Image
 import torch
+import numpy as np
 from torchvision import transforms
 from ultralytics import YOLO
 
@@ -50,7 +51,7 @@ def main():
         st.image(image, caption='Original Image', use_column_width=True)
         
         output = model.predict(image)
-        output_image = Image.fromarray(output)
+        output_image = Image.fromarray(np.array(output))
         st.image(output_image, caption='Processed Image', use_column_width=True)
 
 if __name__ == '__main__':