Coverage for scripts/annotations/coco_txt.py: 75%

1import os 

2import json 

3import numpy as np 

4import argparse 

5 

6def coco_to_txt(coco_json_path, output_txt_dir): 

7 

8 if not os.path.isfile(coco_json_path): 

9 print(f"Error: {coco_json_path} is not a valid file") 

10 return 

11 if not os.path.isdir(output_txt_dir): 

12 os.makedirs(output_txt_dir, exist_ok=True) 

13 

14 # Load COCO JSON 

15 try: 

16 with open(coco_json_path, 'r') as f: 

17 coco_data = json.load(f) 

18 except Exception as e: 

19 print(f"Error loading {coco_json_path}: {str(e)}") 

20 return 

21 

22 # Map image IDs to image info 

23 images = {img['id']: img for img in coco_data['images']} 

24 print(f"Found {len(images)} images in COCO JSON") 

25 

26 # Group annotations by image 

27 annotations_by_image = {} 

28 for ann in coco_data['annotations']: 

29 image_id = ann['image_id'] 

30 if image_id not in annotations_by_image: 

31 annotations_by_image[image_id] = [] 

32 annotations_by_image[image_id].append(ann) 

33 

34 print(f"Found annotations for {len(annotations_by_image)} images") 

35 

36 # Process each image 

37 processed_images = 0 

38 total_txts = 0 

39 

40 for image_id, img_info in images.items(): 

41 try: 

42 filename = img_info['file_name'] 

43 width, height = img_info['width'], img_info['height'] 

44 base_name = os.path.splitext(filename)[0] 

45 

46 # Validate image size 

47 if width != height or width not in [320, 416]: 

48 print(f"Warning: Image {filename} size ({width}x{height}) is not 320x320 or 416x416") 

49 

50 # Get annotations 

51 annotations = annotations_by_image.get(image_id, []) 

52 if not annotations: 

53 print(f"No annotations for {filename}") 

54 continue 

55 

56 # Collect segmentation lines 

57 seg_lines = [] 

58 class_ids_used = set() 

59 for ann in annotations: 

60 cat_id = ann['category_id'] 

61 

62 if 'segmentation' in ann and ann['segmentation']: 

63 # Existing polygon processing (same as your code) 

64 polygon = ann['segmentation'][0] 

65 if len(polygon) < 4 or len(polygon) % 2 != 0: 

66 print(f"Invalid polygon for annotation ID {ann['id']} in {filename}: {polygon}") 

67 continue 

68 

69 normalized_coords = [] 

70 for i in range(0, len(polygon), 2): 

71 x = polygon[i] / width 

72 y = polygon[i + 1] / height 

73 normalized_coords.extend([np.clip(x, 0, 1), np.clip(y, 0, 1)]) 

74 

75 seg_line = [cat_id] + normalized_coords 

76 

77 elif 'bbox' in ann and ann['bbox']: 

78 # bbox format in COCO is [x_min, y_min, width, height] 

79 x_min, y_min, w_box, h_box = ann['bbox'] 

80 x_max = x_min + w_box 

81 y_max = y_min + h_box 

82 

83 # normalize and create polygon as box corners (clockwise) 

84 normalized_coords = [ 

85 np.clip(x_min / width, 0, 1), np.clip(y_min / height, 0, 1), 

86 np.clip(x_max / width, 0, 1), np.clip(y_min / height, 0, 1), 

87 np.clip(x_max / width, 0, 1), np.clip(y_max / height, 0, 1), 

88 np.clip(x_min / width, 0, 1), np.clip(y_max / height, 0, 1) 

89 ] 

90 seg_line = [cat_id] + normalized_coords 

91 

92 seg_lines.append(seg_line) 

93 class_ids_used.add(cat_id) 

94 # Save TXT file if there are annotations 

95 if seg_lines: 

96 txt_path = os.path.join(output_txt_dir, f"{base_name}.txt") 

97 with open(txt_path, 'w') as f: 

98 for line in seg_lines: 

99 # Format: class_id x1 y1 x2 y2 ... xn yn 

100 f.write(' '.join(map(str, line)) + '\n') 

101 print(f"Saved segmentation TXT: {txt_path}") 

102 print(f"Class IDs in {base_name}: {sorted(list(class_ids_used))}") 

103 total_txts += 1 

104 else: 

105 print(f"Skipped saving empty TXT for {filename}") 

106 

107 processed_images += 1 

108 

109 except Exception as e: 

110 print(f"Error processing image ID {image_id} ({filename}): {str(e)}") 

111 continue 

112 

113 print(f"Processed {processed_images} images, saved {total_txts} TXT files") 

114 

115 

116# coco_json_path = '../../8080/train' 

117# output_dir = '../../8080/labels' 

118 

119# for file in os.listdir(coco_json_path): 

120# if file.endswith('.json'): 

121# coco_json_path = os.path.join(coco_json_path, file) 

122# print(f"Processing COCO JSON: {coco_json_path}") 

123 

124# coco_to_txt(coco_json_path, output_dir)