Feature/inference

.gitignore

@@ -0,0 +1,7 @@
+__pycache__
+result/
+
+#*
+.#*
+*~
+*#

README.md

@@ -3,3 +3,22 @@ Sample code and data for Medium post on https://medium.com/fullstackai/how-to-tr
 
 Pytorch version: 1.2.0
 Pycocotools: 2.0.0
+
+# training
+## sample
+$ python train.py
+
+## train on your own data
+- prepare train image data 
+- prepare coco formatted annotation file
+- edit config.py "train_data_dir, train_coco" parameters
+- $ python train.py
+
+# testing
+## sample
+$ python inference.py
+
+## infer on your own data
+- prepare test data
+- edit config.py "test_data_dir, test_img_format"
+- $ python inference.py

config.py

@@ -1,6 +1,6 @@
 # path to your own data and coco file
-train_data_dir = "my_data/train"
-train_coco = "my_data/my_train_coco.json"
+train_data_dir = "data/train"
+train_coco = "data/my_train_coco.json"
 
 # Batch size
 train_batch_size = 1
@@ -13,8 +13,16 @@
 
 # Two classes; Only target class or background
 num_classes = 2
-num_epochs = 10
+num_epochs = 100
 
 lr = 0.005
 momentum = 0.9
 weight_decay = 0.005
+
+# for inference
+save_model_name = 'result/last_model.pth'
+result_img_dir = 'result/imgs'
+detection_threshold = 0.2
+test_data_dir = 'data/test'
+test_img_format = 'jpg'
+

inference.py

@@ -0,0 +1,82 @@
+import torch
+import config
+import utils
+import glob
+import numpy as np
+import cv2
+from pathlib import Path
+
+def draw_bboxes(img, preds, thre, class_colors, save_fname):
+    preds = [{k: v.to('cpu') for k,v in t.items()} for t in preds]
+
+    if len(preds[0]['boxes']) != 0:
+        boxes = preds[0]['boxes'].data.numpy()
+        scores = preds[0]['scores'].data.numpy()
+        print(f"boxes={boxes}, scores = {scores}")
+
+        boxes = boxes[scores >= thre].astype(np.int32)
+        pred_classes = [i for i in preds[0]['labels'].cpu().numpy() ]
+
+        for j, box in enumerate(boxes):
+            color = class_colors[pred_classes[j]]
+            cv2.rectangle(img,
+                        (int(box[0]), int(box[1])),
+                        (int(box[2]), int(box[3])),
+                        color, 2)
+        cv2.imshow('prediction', img)
+        cv2.waitKey(1)
+
+        # save the image
+        cv2.imwrite(save_fname, img)
+
+
+def inference_1img(model, img_name, device, thre, class_colors):
+    in_img = cv2.imread(img_name)
+
+    # display
+    cv2.imshow("input image", in_img)
+    cv2.waitKey(1)
+
+    # convert to tensor
+    img = cv2.cvtColor(in_img, cv2.COLOR_BGR2RGB).astype(np.float32)
+    img /= 255.0
+    img = np.transpose(img, (2,0,1)) # HWC -> CHW
+    img = torch.tensor(img, dtype=torch.float).to(device)
+    img = torch.unsqueeze(img,0) # add batch dim
+
+    # run inference
+    with torch.no_grad():
+        preds = model(img)
+    print(f"inference on {img_name} done.")
+
+    save_fname = str(Path(config.result_img_dir) / Path(img_name).name)
+    draw_bboxes(in_img, preds, thre, class_colors, save_fname)
+
+
+def main():
+    Path(config.result_img_dir).mkdir(parents=True, exist_ok=True)
+
+    # load model
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    saved_name = config.save_model_name
+    checkpoint = torch.load(saved_name, map_location=device)
+    model = utils.get_model_object_detector(config.num_classes)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.to(device).eval()
+
+    # load data
+    test_dir = config.test_data_dir
+    img_format = config.test_img_format
+    test_imgs = glob.glob(f"{test_dir}/*.{img_format}")
+
+    # prepare for drawing
+    class_colors = np.random.uniform(0, 255, size=(config.num_classes, 3))
+
+    # inference
+    for i in range(len(test_imgs)):
+        img_name = test_imgs[i]
+        inference_1img(model, img_name, device, config.detection_threshold, class_colors)
+
+
+if __name__ == "__main__":
+    main()

main.py

@@ -1,11 +1,13 @@
 import torch
 import config
 from utils import (
-    get_model_instance_segmentation,
+    get_model_object_detector,
     collate_fn,
     get_transform,
     myOwnDataset,
+    save_model,
 )
+from pathlib import Path
 
 print("Torch version:", torch.__version__)
 
@@ -34,7 +36,7 @@
     print(annotations)
 
 
-model = get_model_instance_segmentation(config.num_classes)
+model = get_model_object_detector(config.num_classes)
 
 # move model to the right device
 model.to(device)
@@ -47,6 +49,9 @@
 
 len_dataloader = len(data_loader)
 
+# crete output directory
+Path("result/").mkdir(parents=True, exist_ok=True)
+
 # Training
 for epoch in range(config.num_epochs):
     print(f"Epoch: {epoch}/{config.num_epochs}")
@@ -64,3 +69,5 @@
         optimizer.step()
 
         print(f"Iteration: {i}/{len_dataloader}, Loss: {losses}")
+
+save_model(config.num_epochs, model, optimizer, save_name=config.save_model_name)

train.py

@@ -0,0 +1,73 @@
+import torch
+import config
+from utils import (
+    get_model_object_detector,
+    collate_fn,
+    get_transform,
+    myOwnDataset,
+    save_model,
+)
+from pathlib import Path
+
+print("Torch version:", torch.__version__)
+
+# create own Dataset
+my_dataset = myOwnDataset(
+    root=config.train_data_dir, annotation=config.train_coco, transforms=get_transform()
+)
+
+# own DataLoader
+data_loader = torch.utils.data.DataLoader(
+    my_dataset,
+    batch_size=config.train_batch_size,
+    shuffle=config.train_shuffle_dl,
+    num_workers=config.num_workers_dl,
+    collate_fn=collate_fn,
+)
+
+
+# select device (whether GPU or CPU)
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+# DataLoader is iterable over Dataset
+for imgs, annotations in data_loader:
+    imgs = list(img.to(device) for img in imgs)
+    annotations = [{k: v.to(device) for k, v in t.items()} for t in annotations]
+    print(annotations)
+
+
+model = get_model_object_detector(config.num_classes)
+
+# move model to the right device
+model.to(device)
+
+# parameters
+params = [p for p in model.parameters() if p.requires_grad]
+optimizer = torch.optim.SGD(
+    params, lr=config.lr, momentum=config.momentum, weight_decay=config.weight_decay
+)
+
+len_dataloader = len(data_loader)
+
+# crete output directory
+Path("result/").mkdir(parents=True, exist_ok=True)
+
+# Training
+for epoch in range(config.num_epochs):
+    print(f"Epoch: {epoch}/{config.num_epochs}")
+    model.train()
+    i = 0
+    for imgs, annotations in data_loader:
+        i += 1
+        imgs = list(img.to(device) for img in imgs)
+        annotations = [{k: v.to(device) for k, v in t.items()} for t in annotations]
+        loss_dict = model(imgs, annotations)
+        losses = sum(loss for loss in loss_dict.values())
+
+        optimizer.zero_grad()
+        losses.backward()
+        optimizer.step()
+
+        print(f"Iteration: {i}/{len_dataloader}, Loss: {losses}")
+
+save_model(config.num_epochs, model, optimizer)

utils.py

@@ -82,7 +82,7 @@ def collate_fn(batch):
     return tuple(zip(*batch))
 
 
-def get_model_instance_segmentation(num_classes):
+def get_model_object_detector(num_classes):
     # load an instance segmentation model pre-trained pre-trained on COCO
     model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)
     # get number of input features for the classifier
@@ -91,3 +91,13 @@ def get_model_instance_segmentation(num_classes):
     model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
 
     return model
+
+def save_model(epoch, model, optimizer):
+    """
+    Function to save the trained model till current epoch, or whenver called
+    """
+    torch.save({
+                'epoch': epoch+1,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                }, 'result/last_model.pth')