Dede16 commited on
Commit
4bb4d2c
·
verified ·
1 Parent(s): e35cc54

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +50 -0
app.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import uuid
3
+ import pandas as pd
4
+ import cv2
5
+ from ultralytics import YOLO
6
+
7
+ # Initialize the YOLO model
8
+ model = YOLO('runs/segment/train/weights/best.pt')
9
+
10
+ # Define directories
11
+ image_dir = 'datasetsw/valid/ripe'
12
+ output_dir = 'datasetsw/valid/test_ripe'
13
+
14
+ # List all image files in the directory
15
+ image_files = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith(('.jpg', '.jpeg', '.png'))]
16
+
17
+ # Initialize an empty list to store the results
18
+ results_list = []
19
+
20
+ # Process each image
21
+ for image_file in image_files:
22
+ results = model(image_file)
23
+ annotated_frame = results[0].plot()
24
+
25
+ total_objects = len(results[0].boxes)
26
+ labels = results[0].boxes.cls.tolist()
27
+ matang_count = labels.count(0)
28
+ mentah_count = labels.count(1)
29
+
30
+ # Generate a unique filename and save the annotated image
31
+ unique_id = str(uuid.uuid4())
32
+ combined_filename = f"{unique_id}.jpg"
33
+ output_path = os.path.join(output_dir, combined_filename)
34
+ cv2.imwrite(output_path, annotated_frame)
35
+
36
+ # Append the results to the list
37
+ results_list.append({
38
+ 'image_file': image_file,
39
+ 'total_objects': total_objects,
40
+ 'matang_count': matang_count,
41
+ 'mentah_count': mentah_count
42
+ })
43
+
44
+ # Convert the results list to a dataframe
45
+ results_df = pd.DataFrame(results_list)
46
+
47
+ # Save the dataframe to a CSV file
48
+ results_df.to_csv(os.path.join(output_dir, 'detection_results.csv'), index=False)
49
+
50
+ cv2.destroyAllWindows()