Coverage for scripts/annotations/count_labels.py: 70%

1import os 

2from collections import Counter 

3import shutil 

4 

5def count_class_ids(label_dir): 

6 if not os.path.exists(label_dir): 

7 print(f"Label directory does not exist: {label_dir}") 

8 return {} 

9 

10 # Initialize counter for class IDs 

11 class_counts = Counter() 

12 

13 # Get list of txt files 

14 txt_files = [f for f in os.listdir(label_dir) if f.endswith('.txt')] 

15 

16 if not txt_files: 

17 print(f"No .txt files found in {label_dir}") 

18 return {} 

19 

20 # Process each annotation file 

21 for txt_file in txt_files: 

22 file_path = os.path.join(label_dir, txt_file) 

23 try: 

24 with open(file_path, 'r') as f: 

25 lines = f.readlines() 

26 

27 # Count class IDs in each line 

28 for line in lines: 

29 parts = line.strip().split() 

30 if parts: # Check if line is not empty 

31 try: 

32 class_id = parts[0] # Class ID is the first element 

33 # if class_id == '10' or class_id == '11': 

34 # print(f"class {class_id} in {txt_file}") 

35 class_counts[class_id] += 1 

36 except (IndexError, ValueError): 

37 print(f"Invalid line in {txt_file}: {line.strip()}") 

38 continue 

39 except OSError as e: 

40 print(f"Error reading file {txt_file}: {e}") 

41 continue 

42 

43 # Print results 

44 if class_counts: 

45 print("\nClass ID Counts:") 

46 for class_id, count in sorted(class_counts.items()): 

47 print(f"Class ID {class_id}: {count} occurrences") 

48 print(f"Total annotations: {sum(class_counts.values())}") 

49 print(f"Total files processed: {len(txt_files)}") 

50 else: 

51 print("No valid annotations found.") 

52 

53 return dict(class_counts) 

54 

55# Example usage 

56if __name__ == "__main__": 

57 label_directory = "../../seame_training/train/labels" # Replace with your label directory path 

58 class_counts = count_class_ids(label_directory) 

59 # label_directory = "../../new_seame/val/labels" # Replace with your label directory path 

60 # class_counts = count_class_ids(label_directory) 

61 # label_directory = "../dataset/labels/train" # Ensure absolute path 

62 # class_counts = count_class_ids(label_directory)