Coverage for scripts/annotations/remap_classid.py: 88%

1import os 

2 

3def change_annotation_labels(label_dir, label_map): 

4 

5 

6 # Get list of txt files in label directory 

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

8 

9 # Counter for modified files 

10 modified_count = 0 

11 

12 for txt_file in txt_files: 

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

14 modified = False 

15 new_lines = [] 

16 

17 # Read the file 

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

19 lines = f.readlines() 

20 

21 # Process each line 

22 for line in lines: 

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

24 if parts and parts[0] in label_map: # Check if line is not empty 

25 # Check if the first element (label) needs to be changed 

26 parts[0] = label_map[parts[0]] 

27 modified = True 

28 new_lines.append(' '.join(parts) + '\n') 

29 else: 

30 new_lines.append(line) 

31 

32 # Write back to file only if modified 

33 if modified: 

34 with open(file_path, 'w') as f: 

35 f.writelines(new_lines) 

36 modified_count += 1 

37 

38 print(f"Total files modified: {modified_count}") 

39 

40if __name__ == "__main__": 

41 label_map = { 

42 '2': '5', 

43 # Add other mappings if needed 

44 } 

45 output_directory = "../../8080/labels/" # Adjust this path as needed 

46 change_annotation_labels(output_directory, label_map) 

47