yacwc

data.py

@@ -1,18 +1,19 @@
 # %%
+import os
+import numpy as np
+import torch
+from PIL import Image
 import torchvision
 from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
-from collections import defaultdict as ddict
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
 import json
 import torch
-from torchvision import datasets, transforms as T
-import cv2
+from torchvision import transforms as T
 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
@@ -28,9 +29,6 @@ def get_transform(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')
@@ -43,12 +41,12 @@ def create_map(list_in, from_key, to_key):
 
 
 class iNaturalistDataset(torch.utils.data.Dataset):
-    def __init__(self, validation=False, train=False):
+    def __init__(self, validation=False, train=False, transforms = None):
 
         self.validation = validation
         self.train = train
+        self.transforms = transforms
 
-        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")
@@ -65,7 +63,7 @@ class iNaturalistDataset(torch.utils.data.Dataset):
 
         for category in f["categories"]:
             if category["supercategory"] == "Aves":
-                if category['name'] in ['Archilochus colubris','Icterus galbula']:
+                if category['name'] in ['Archilochus colubris']:#,'Icterus galbula']:
                     print(category['name'])
                     categories.append(category)
 
@@ -101,44 +99,35 @@ class iNaturalistDataset(torch.utils.data.Dataset):
         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)
+        img_path = c_image["path"]
+        img = Image.open(img_path).convert("RGB")
+    
         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)
+        boxes = bbox
         target = dict()
-        target["boxes"] = boxes
+        target["boxes"] = torch.as_tensor([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['area'] =  torch.as_tensor([annot['area']])
         target['iscrowd'] = torch.zeros((1,), dtype=torch.int64)
         
-        return image, target
+        
+        if self.transforms is not None:
+            img, target = self.transforms(img, target)
+
+        return img, target
+# %%
+# v = iNaturalistDataset(validation=True)
+    
 
 # v = iNaturalistDataset(validation= True)
 # o = v[10]
@@ -149,24 +138,45 @@ class iNaturalistDataset(torch.utils.data.Dataset):
 # plt.imshow(ox.permute([1,2,0]))
 # plt.savefig('crap2.png')
 
+def get_model(num_classes):
+    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
+    num_classes = 2  # 1 class (person) + background
+    in_features = model.roi_heads.box_predictor.cls_score.in_features
+    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
+    return model
+
+
+import transforms as T
+
+def get_transform(train):
+    transforms = []
+    transforms.append(T.ToTensor())
+    if train:
+        transforms.append(T.RandomHorizontalFlip(0.5))
+    return T.Compose(transforms)
+
+from engine import train_one_epoch, evaluate
+import utils
 # %%
 def run():
-    val_dataset = iNaturalistDataset(validation=True)
-    train_dataset = iNaturalistDataset(train=True)
+    val_dataset = iNaturalistDataset(validation=True, transforms = get_transform(train=True))
+    train_dataset = iNaturalistDataset(train=True, transforms = get_transform(train=False))
 
 
     train_data_loader = torch.utils.data.DataLoader(
-        train_dataset, batch_size=16, shuffle=True, num_workers=4,  collate_fn=utils.collate_fn
+        train_dataset, batch_size=8, shuffle=True, num_workers=1,  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
+        val_dataset, batch_size=8, shuffle=True, num_workers=1,  collate_fn=utils.collate_fn
     )
 
+    import torchvision
+    from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+    num_classes = 2
+
+ 
+    model = get_model(num_classes)
     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,