Update sample.py

sample.py

@@ -13,7 +13,7 @@ def load_model():
     return model
 
 # Define YOLO processing function
-def process_image(image, model):
+def process_image(image):
     # Preprocess the image
     preprocess = transforms.Compose([
         transforms.Resize((416, 416)),
@@ -23,17 +23,17 @@ def process_image(image, model):
     input_tensor = preprocess(image)
     input_batch = input_tensor.unsqueeze(0)
     
-    # Perform inference
-    with torch.no_grad():
-        output = model(input_batch)
+    # # 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
+    # # 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())
+    # # Convert tensor to PIL Image
+    # output_image = transforms.ToPILImage()(output[0].cpu().squeeze())
 
-    return output_image
+    return input_batch
 
 # Main Streamlit code
 def main():
@@ -50,7 +50,7 @@ def main():
         image = Image.open(uploaded_file)
         st.image(image, caption='Original Image', use_column_width=True)
         
-        output = model.predict(image)
+        output = process_image(image)
         output_image = Image.fromarray(np.array(output))
         st.image(output_image, caption='Processed Image', use_column_width=True)