# %% import torchvision from torchvision.models.detection.faster_rcnn import FastRCNNPredictor from collections import defaultdict as ddict import json import torch from torchvision import datasets, transforms as T import cv2 import numpy as np import os import sys sys.path.append(r"K:\Designs\ML\inaturalist_models\data_aug") sys.path.append(r"K:\Designs\ML\inaturalist_models\vision") from references.detection import utils, engine import data_aug import bbox_util def get_transform(train): transforms = [] transforms.append(T.ToTensor()) if train: transforms.append(T.RandomHorizontalFlip(0.5)) return T.Compose(transforms) IMAGE_MEAN = [0.485, 0.456, 0.406] IMAGE_STD = [0.229, 0.224, 0.225] PATH_ROOT = r"D:\ishan\ml\inaturalist\\" device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') def create_map(list_in, from_key, to_key): cmap = dict() for l in list_in: cmap[l[from_key]] = l[to_key] return cmap class iNaturalistDataset(torch.utils.data.Dataset): def __init__(self, validation=False, train=False): self.validation = validation self.train = train self.transforms = T.Compose([T.Resize(600, max_size=1024), T.ToTensor()]) if validation: json_path = os.path.join(PATH_ROOT, r"val_2017_bboxes\val_2017_bboxes.json") elif train: json_path = os.path.join( PATH_ROOT, r"train_2017_bboxes\train_2017_bboxes.json" ) with open(json_path, "r") as rj: f = json.load(rj) categories = list() image_info = dict() for category in f["categories"]: if category["supercategory"] == "Aves": if category['name'] in ['Archilochus colubris','Icterus galbula']: print(category['name']) categories.append(category) categories = sorted(categories, key=lambda k: k["name"]) for idx, cat in enumerate(categories): cat["new_id"] = idx + 1 orig_to_new_id = create_map(categories, "id", "new_id") for annot in f["annotations"]: if annot["category_id"] in orig_to_new_id: annot["new_category_id"] = orig_to_new_id[annot["category_id"]] id = annot["image_id"] if id not in image_info: image_info[id] = dict() annot["bbox"][2] += annot["bbox"][0] annot["bbox"][3] += annot["bbox"][1] image_info[id]["annotation"] = annot for img in f["images"]: id = img["id"] path = os.path.join(PATH_ROOT, img["file_name"]) height = img["height"] width = img["width"] if id in image_info: image_info[id].update({"path": path, "height": height, "width": width}) for idx, (id, im_in) in enumerate(image_info.items()): im_in["idx"] = idx self.images = image_info self.categories = categories self.idx_to_id = [x for x in self.images] self.num_classes = len(self.categories) + 1 self.num_samples = len(self.images) self.transforms = [ data_aug.RandomHorizontalFlip(0.5), data_aug.Resize(600), ] self.pre_transform = T.Compose([T.ToTensor()])#],T.Normalize(mean=[0.485, 0.456, 0.406], #std=[0.229, 0.224, 0.225])]) def __len__(self): return self.num_samples def transform(self, img, bbox): for x in self.transforms: img, bbox = x(img, bbox) img = self.pre_transform(img) return img, bbox def __getitem__(self, idx): idd = self.idx_to_id[idx] c_image = self.images[idd] # print(c_image, idx, self.validation, self.train) # breakpoint() image = np.asarray(cv2.imread(c_image["path"])[:,:,::-1].copy(),dtype=np.float32) annot = c_image["annotation"] bbox = annot["bbox"] bbox.append(annot["new_category_id"]) bbox = np.asarray([bbox], dtype=np.float32) image, bbox = self.transform(image.copy(), bbox.copy()) boxes = torch.as_tensor(bbox[:,:4], dtype=torch.float32) target = dict() target["boxes"] = boxes target["labels"] = torch.as_tensor([annot["new_category_id"]], dtype=torch.int64) target['image_id'] = torch.tensor([annot['image_id']]) target['area'] = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0]) target['iscrowd'] = torch.zeros((1,), dtype=torch.int64) return image, target # v = iNaturalistDataset(validation= True) # o = v[10] # %% # oimage = t.tensor(o[0]*255, dtype=t.uint8) # import matplotlib.pyplot as plt # ox = draw_bounding_boxes(oimage, o[1]['boxes'], width=1) # plt.imshow(ox.permute([1,2,0])) # plt.savefig('crap2.png') # %% def run(): val_dataset = iNaturalistDataset(validation=True) train_dataset = iNaturalistDataset(train=True) train_data_loader = torch.utils.data.DataLoader( train_dataset, batch_size=16, shuffle=True, num_workers=4, collate_fn=utils.collate_fn ) val_data_loader = torch.utils.data.DataLoader( val_dataset, batch_size=16, shuffle=True, num_workers=4, collate_fn=utils.collate_fn ) model = torchvision.models.detection.fasterrcnn_resnet50_fpn( pretrained=True, num_classes=train_dataset.num_classes, progress=True ) model.to(device) # construct an optimizer params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005) # and a learning rate scheduler lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1) # let's train it for 10 epochs num_epochs = 10 for epoch in range(num_epochs): print(epoch) torch.save(model.state_dict(), 'model_weights_start_'+str(epoch)+ '.pth') # train for one epoch, printing every 10 iterations engine.train_one_epoch(model, optimizer, train_data_loader, device, epoch, print_freq=10) torch.save(model.state_dict(), 'model_weights_post_train_'+str(epoch)+ '.pth') # update the learning rate lr_scheduler.step() torch.save(model.state_dict(), 'model_weights_post_step_'+str(epoch)+ '.pth') # evaluate on the test dataset engine.evaluate(model, val_data_loader, device=device) if __name__ == "__main__": run() # # %% # json_path = os.path.join( # PATH_ROOT, r"train_2017_bboxes\train_2017_bboxes.json" # ) # with open(json_path, "r") as rj: # f = json.load(rj) # # %% # image_id: 2358