Update sample.py

sample.py

@@ -1,60 +1,57 @@
 import streamlit as st
 from PIL import Image
-from io import BytesIO
 import torch
-from pathlib import Path
-from PIL import Image
-from torchvision.transforms import functional as F
-from utils.plots import plot_one_box
-
-# Import YOLOv5
-from models.experimental import attempt_load
-from utils.general import non_max_suppression, scale_coords
-from utils.torch_utils import select_device
-
-# Function to load YOLOv5 model
-@st.cache(allow_output_mutation=True)
-def load_model(weights='best.pt', device=''):
-    device = select_device(device)
-    model = attempt_load(weights, map_location=device)
+from torchvision import transforms
+from ultralytics import YOLO
+
+# Load the YOLO model
+@st.cache
+def load_model():
+    # Replace 'model.pt' with the path to your YOLO model file
+    model = YOLO('best.pt')
     return model
 
-# Function to process image through YOLOv5
+# Define YOLO processing function
 def process_image(image, model):
-    # Convert PIL image to Torch tensor
-    img_tensor = F.to_tensor(image)
-    img_tensor = img_tensor.unsqueeze(0)
-
-    # Inference
-    results = model(img_tensor)
-    results = non_max_suppression(results, conf_thres=0.4, iou_thres=0.5)
-
-    # Draw bounding boxes on the image
-    for i, det in enumerate(results):
-        if det is not None and len(det):
-            det[:, :4] = scale_coords(img_tensor.shape[2:], det[:, :4], image.size)
-            for *xyxy, conf, cls in reversed(det):
-                plot_one_box(xyxy, image, label=f'{model.names[int(cls)]} {conf:.2f}', color='red', line_thickness=3)
+    # Preprocess the image
+    preprocess = transforms.Compose([
+        transforms.Resize((416, 416)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    input_tensor = preprocess(image)
+    input_batch = input_tensor.unsqueeze(0)
+    
+    # Perform inference
+    with torch.no_grad():
+        output = model(input_batch)
+    
+    # Post-process the output (e.g., draw bounding boxes)
+    # Replace this with your post-processing code
+    
+    # Convert tensor to PIL Image
+    output_image = transforms.ToPILImage()(output[0].cpu().squeeze())
 
-    return image
+    return output_image
 
-# Main function
+# Main Streamlit code
 def main():
-    st.title('YOLOv5 Image Detection')
+    st.title('YOLO Image Detection')
     
     # Upload image file
     uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
+    
     if uploaded_file is not None:
-        # Load YOLOv5 model
+        # Load YOLO model
         model = load_model()
-
+        
         # Process uploaded image
         image = Image.open(uploaded_file)
         st.image(image, caption='Original Image', use_column_width=True)
-
-        output_image = process_image(image, model)
+        
+        output = model.predict(image)
+        output_image = Image.fromarray(output)
         st.image(output_image, caption='Processed Image', use_column_width=True)
 
 if __name__ == '__main__':
-    main()
+    main()