Update sample.py

sample.py

@@ -51,18 +51,20 @@ def main():
         # Process uploaded image
         # image_l = Image.open(uploaded_file)
         image = plt.imread(uploaded_file)
+        # Process uploaded image
+        # image = Image.open(uploaded_file)
         st.image(image, caption='Original Image', use_column_width=True)
         
         # output_l = process_image(image)
         # output = np.asarray(output_l)
         # output_image = Image.fromarray(output)
         results = model(image)
-        for result in results:
-            boxes = result['boxes']  # Assuming 'boxes' is a list of bounding box coordinates [x_min, y_min, x_max, y_max]
-            for box in boxes:
-                x_min, y_min, x_max, y_max = box  # Extract coordinates
-                cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
+        for result in results.pred[0]:
+            box = result.xyxy[0].tolist()  # Extract bounding box coordinates as a list [x_min, y_min, x_max, y_max]
+            x_min, y_min, x_max, y_max = box
+            cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
         st.image(image, caption='Processed Image', use_column_width=True)
 
+
 if __name__ == '__main__':
     main()