Oxytocin neurons signal state-dependent transitions from rest to thermogenesis and behavioral arousal.

1. import os
2. import torch
3. import torch.nn as nn
4. import torchvision
5. import json
6. from torchvision.models import resnet50, ResNet50_Weights
7. from torchvision.models.detection import MaskRCNN
8. from torchvision.models.detection.rpn import AnchorGenerator
9. from src.data.Data_alignment import prepare_dataloader
10. from src.utils.utils import keypoints_to_heatmaps, visualize_prediction_images, filter_duplicate_classes
11. import csv
13. os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" # Ignore
15. # Define directories
16. csv_directory = r'F:\DLC_thermography-JLanden-2024-07-09\videos' #Keypoints from DLC
17. image_directory = r'F:\DLC_thermography-JLanden-2024-07-09\labeled-data' # Images
18. segments_dirs = r'F:\Segmentation\training_segments' # Training segments from Darwin
19. save_to_dir = r'F:\Segmentation\compare_02-13-25_BS8_check' #Where output is saved
21. # Parameters
22. num_epochs = 200
23. validate_every_n_epochs = 200
24. batch_size = 8 # Adjust as needed
25. val_split = 0.2 # 20% of the data for validation
26. #mask_threshold = 0.5 # 0.5 is normal
27. mask_display_threshold = 0.9
28. save_weights = False
29. save_losses = True
31. image_size = (348, 464) # H, W
32. num_keypoints = 19
33. num_classes = 4 # Set this according to your dataset (# classes + background)
34. class_names = ['background', 'implanted_mouse', 'implanted_bat', 'implanted_rump' #,
35. # , 'extra_mouse', 'extra_bat', 'extra_rump'
36. ]
38. loss_history = [] # Store loss values for CSV
39. colors = [(1, 0, 0), (.1, .6, .1), (0, 0, 1)] #R, G, B: Must match number of class names
41. # Attention Mechanism
42. class AttentionLayer(nn.Module):
43. def __init__(self, img_channels, keypoint_channels, attention_channels):
44. super(AttentionLayer, self).__init__()
45. # Convolution layers for image and keypoints (input)
46. self.img_conv1 = nn.Conv2d(img_channels, attention_channels, kernel_size=1)
47. self.img_conv2 = nn.Conv2d(attention_channels, attention_channels, kernel_size=3, padding=1)
48. self.img_conv3 = nn.Conv2d(attention_channels, attention_channels, kernel_size=3, padding=1)
50. self.keypoint_conv1 = nn.Conv2d(keypoint_channels, attention_channels, kernel_size=1)
51. self.keypoint_conv2 = nn.Conv2d(attention_channels, attention_channels, kernel_size=3, padding=1)
52. self.keypoint_conv3 = nn.Conv2d(attention_channels, attention_channels, kernel_size=3, padding=1)
54. # Activation function
55. self.relu = nn.ReLU()
57. # Convolution for the attention mechanism (combine image and keypoint features)
58. self.attention_conv1 = nn.Conv2d(attention_channels, attention_channels, kernel_size=3, padding=1)
59. self.attention_conv2 = nn.Conv2d(attention_channels, attention_channels, kernel_size=3, padding=1)
60. self.sigmoid = nn.Sigmoid()
62. def forward(self, images, keypoints):
63. # Process the image through multiple convolutions
64. img_out = self.relu(self.img_conv1(images))
65. img_out = self.relu(self.img_conv2(img_out))
66. img_out = self.relu(self.img_conv3(img_out))
68. # Process the keypoints through multiple convolutions
69. keypoint_out = self.relu(self.keypoint_conv1(keypoints))
70. keypoint_out = self.relu(self.keypoint_conv2(keypoint_out))
71. keypoint_out = self.relu(self.keypoint_conv3(keypoint_out))
73. # Combine image and keypoint features using element-wise multiplication
74. combined = img_out * keypoint_out
76. # Apply further attention convolutions
77. attention_map = self.relu(self.attention_conv1(combined))
78. attention_map = self.sigmoid(self.attention_conv2(attention_map))
80. # Apply attention to image features
81. enhanced_img_features = images * attention_map
82. return enhanced_img_features
85. # Custom Mask R-CNN Model with Attention
86. class SkeletonGuidedMaskRCNN(nn.Module):
87. def __init__(self, num_classes):
88. super(SkeletonGuidedMaskRCNN, self).__init__()
90. # Attention layer first to overlay images with keypoints
91. self.attention = AttentionLayer(img_channels=3, keypoint_channels=num_keypoints, attention_channels=3)
93. # Load a pre-trained Mask R-CNN backbone
94. backbone = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
95. backbone = nn.Sequential(*list(backbone.children())[:-2]) # Remove last fully connected layer
97. self.backbone = backbone
98. self.backbone.out_channels = 2048
100. # Create an Anchor Generator for the FPN (feature pyramid network)
101. anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),), aspect_ratios=((0.5, 1.0, 2.0),) * 5)
103. # Define RoI Align and FPN
104. roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], output_size=4, sampling_ratio=2)
106. # Define Mask R-CNN model with the modified backbone
107. self.model = MaskRCNN(self.backbone, num_classes=num_classes, rpn_anchor_generator=anchor_generator,
108. box_roi_pool=roi_pooler, mask_roi_pool=roi_pooler)
110. def forward(self, images, keypoints, targets=None, return_predictions=False):
111. # Apply attention to images using keypoints
112. enhanced_img_features = self.attention(images, keypoints)
114. # Use enhanced image features in the Mask R-CNN model
115. if self.training and not return_predictions:
116. return self.model(enhanced_img_features, targets)
117. else: # Validate
118. return self.model(enhanced_img_features)
121. # Training script
122. # Initialize the custom loss
123. # loss_fn = CustomLoss()
125. ########################
126. def save_predictions_to_json(predictions, save_to_dir, videoID, frameID):
127. os.makedirs(f"{save_to_dir}/predictions/{videoID}", exist_ok=True)
128. json_path = os.path.join(save_to_dir, f"predictions", videoID, f"predictions_{frameID}.txt")
130. serializable_predictions = []
131. for pred in predictions:
132. serializable_predictions.append({
133. 'labels': pred['labels'].cpu().tolist(), # Convert tensor to list
134. 'scores': pred['scores'].cpu().tolist(), # Convert tensor to list
135. 'boxes': pred['boxes'].cpu().tolist(), # Convert tensor to list
136. 'masks': pred['masks'].cpu().numpy().tolist() # Convert tensor to nested lists
137. })
139. # Write the predictions to a file in JSON format
140. with open(json_path, 'w') as f:
141. json.dump(serializable_predictions, f, indent=4)
142. #####################
145. def train_model(train_loader, val_loader, num_classes, num_epochs=num_epochs):
146. device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
148. model = SkeletonGuidedMaskRCNN(num_classes=num_classes).to(device)
149. optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
151. for epoch in range(num_epochs):
152. model.train()
153. total_loss = 0
155. for batch_idx, (images, segments, skeletons) in enumerate(train_loader):
156. images = [image.to(device) for image in images]
157. H, W = images[0].shape[1], images[0].shape[2]
158. # Convert images to tensors
159. images = torch.stack([image.to(device) for image in images]).to(device)
161. keypoints = skeletons.to(device)
162. kp_heatmaps = keypoints_to_heatmaps(keypoints, H, W)
163. segments = [{k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in segment.items()} for
164. segment in segments]
166. optimizer.zero_grad()
167. outputs = model(images, kp_heatmaps, segments)
169. # Use the losses from the model's output
170. rcnn_losses = sum(loss for loss in outputs.values())
171. total_loss = rcnn_losses
173. total_loss.backward()
174. optimizer.step()
175. total_loss += total_loss.item()
177. avg_loss = total_loss / len(train_loader)
178. print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {avg_loss}")
179. # Store loss history for CSV
180. loss_history.append([epoch + 1, avg_loss.item()])
182. if (epoch + 1) % validate_every_n_epochs == 0:
183. model.eval()
184. predictions_list = []
185. for batch_idx, (images, segments, skeletons) in enumerate(val_loader):
186. images = [image.to(device) for image in images]
187. H, W = images[0].shape[1], images[0].shape[2]
188. images = torch.stack([image.to(device) for image in images]).to(device)
190. keypoints = skeletons.to(device)
191. kp_heatmaps = keypoints_to_heatmaps(keypoints, H, W)
192. segments = [{k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in segment.items()} for
193. segment in segments]
195. with torch.no_grad(): # Disable gradient calculation
196. # Get predictions instead of losses
197. predictions = model(images, kp_heatmaps)
198. predictions = filter_duplicate_classes(predictions)
199. predictions_list.append(predictions)
201. # Optionally visualize here
202. visualize_prediction_images(
203. images=images,
204. predictions=predictions,
205. epoch=epoch + 1,
206. phase="validate",
207. batch_idx=batch_idx,
208. save_img_dir=save_to_dir,
209. class_names=class_names,
210. colors=colors,
211. keypoints=keypoints,
212. threshold=mask_display_threshold
213. )
215. else:
216. continue
218. # Save model weights
219. if save_weights is True:
220. model_save_path = os.path.join(save_to_dir, f'sg-mrcnn_epochs_{num_epochs}.pth')
221. torch.save(model.state_dict(), model_save_path)
222. print(f"Model weights saved to {model_save_path}")
224. if save_losses is True:
225. loss_csv_path = os.path.join(save_to_dir, f"main_loss_{num_epochs}.csv")
226. # Save loss history to CSV
227. with open(loss_csv_path, mode="w", newline="") as file:
228. writer = csv.writer(file)
229. writer.writerow(["Epoch", "Loss"]) # Write header
230. writer.writerows(loss_history) # Write epoch, loss pairs
232. return model
235. if __name__ == "__main__":
236. # Assuming train_loader and val_loader are already created with your data
237. train_loader, val_loader = prepare_dataloader(csv_directory, segments_dirs, image_directory,
238. batch_size, val_split) # Implement this function
240. trained_model = train_model(train_loader, val_loader, num_classes)

SG_MaskRCNN_train.py at commit 06d57d8, under MIT · at the source