Coverage for scripts/annotations/bbox_seg.py: 80%

1import os 

2 

3def convert_bbox_to_segmentation(label_dir, output_dir): 

4 os.makedirs(output_dir, exist_ok=True) 

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 processed files 

10 processed_count = 0 

11 

12 for txt_file in txt_files: 

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

14 output_path = os.path.join(output_dir, txt_file) 

15 

16 # Read the input annotation file 

17 with open(input_path, 'r') as f: 

18 lines = f.readlines() 

19 

20 new_lines = [] 

21 for line in lines: 

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

23 if len(parts) < 5: # Skip invalid lines 

24 print(f"Skipping invalid line in {txt_file}: {line.strip()}") 

25 new_lines.append(line) 

26 continue 

27 

28 try: 

29 # Parse bounding box annotation 

30 class_id = parts[0] 

31 center_x = float(parts[1]) 

32 center_y = float(parts[2]) 

33 width = float(parts[3]) 

34 height = float(parts[4]) 

35 

36 # Calculate the four corners of the bounding box 

37 half_width = width / 2 

38 half_height = height / 2 

39 x1 = center_x - half_width # Top-left x 

40 y1 = center_y - half_height # Top-left y 

41 x2 = center_x + half_width # Top-right x 

42 y2 = center_y - half_height # Top-right y 

43 x3 = center_x + half_width # Bottom-right x 

44 y3 = center_y + half_height # Bottom-right y 

45 x4 = center_x - half_width # Bottom-left x 

46 y4 = center_y + half_height # Bottom-left y 

47 

48 # Ensure coordinates are within [0, 1] 

49 coords = [x1, y1, x2, y2, x3, y3, x4, y4] 

50 coords = [max(0.0, min(1.0, coord)) for coord in coords] 

51 

52 # Format the new segmentation line 

53 new_line = f"{class_id} {' '.join(map(str, coords))}\n" 

54 new_lines.append(new_line) 

55 

56 except ValueError: 

57 print(f"Error parsing line in {txt_file}: {line.strip()}") 

58 new_lines.append(line) 

59 continue 

60 

61 # Write to output file 

62 with open(output_path, 'w') as f: 

63 f.writelines(new_lines) 

64 processed_count += 1 

65 print(f"Processed: {txt_file}") 

66 

67 print(f"Total files processed: {processed_count}") 

68 

69 

70if __name__ == "__main__": 

71 label_directory = "../../5050/train/labels" # Path to the label directory 

72 output_directory = "../../5050/train/labels_seg" 

73 convert_bbox_to_segmentation(label_directory, output_directory)