diff --git a/.gitignore b/.gitignore index ae3d542..af81301 100644 --- a/.gitignore +++ b/.gitignore @@ -12,9 +12,19 @@ handling/ tests/ scripts/ *.code-workspace - +data/ +*.JPG configs/train/* configs/test/* !configs/*/faster_rcnn.yaml !configs/*/dla34.yaml -!configs/*/herdnet.yaml \ No newline at end of file +!configs/*/herdnet.yaml +wandb +runs +*pt +nohup.out +yolo-runs +tmp +cache.db +template_annotations.json +template_preannotations.json diff --git a/animaloc/data/batch_utils.py b/animaloc/data/batch_utils.py index bcbfd24..1fcefd1 100644 --- a/animaloc/data/batch_utils.py +++ b/animaloc/data/batch_utils.py @@ -36,6 +36,38 @@ def collate_fn(batch): batched_images = cat_list(images) return batched_images , targets +def val_collate_fn(self, batch: tuple) -> tuple[torch.Tensor, dict]: + """collate_fn used to create the validation dataloader + + Args: + batch (tuple): (img:torch.Tensor, targets:dict) + + Returns: + tuple: (image, target) + """ + + batched = dict(points=[], labels=[]) + batch_img = torch.stack([p[0] for p in batch]) + targets = [p[1] for p in batch] + keys = targets[0].keys() + + # get non_empty samples indidces -> set difference + non_empty_idx = [i for i, a in enumerate(targets) if len(a["labels"]) > 0] + targets_empty = [ + targets[i] for i in list(set(range(len(batch))) - set(non_empty_idx)) + ] + targets = [targets[i] for i in non_empty_idx] + + # Creating batch + for k in keys: + batched[k] = [] # initialize to be empty list + if k == "points" or k=='labels': + batched[k] = [a[k].cpu().tolist() for a in targets] + if len(targets_empty) > 0: + batched[k] = batched[k] + [[]] * len(targets_empty) + return batch_img, batched + + def to_xywh(bbox): ''' Bbox from [x_min,y_min,x_max,y_max] to [x,y,width,height] ''' width = bbox[2] - bbox[0] @@ -83,4 +115,4 @@ def show_batch(sample_batched, denormalize=True): ax.add_patch(rect) ax.set_title(f'Batch of {batch_size} images') - ax.axis('off') \ No newline at end of file + ax.axis('off') diff --git a/animaloc/datasets/folder.py b/animaloc/datasets/folder.py index a53cfbf..0eda781 100644 --- a/animaloc/datasets/folder.py +++ b/animaloc/datasets/folder.py @@ -22,7 +22,6 @@ from typing import Optional, List, Any, Dict -from ..data.types import BoundingBox from ..data.utils import group_by_image from .register import DATASETS @@ -57,7 +56,8 @@ def __init__( csv_file: str, root_dir: str, albu_transforms: Optional[list] = None, - end_transforms: Optional[list] = None + end_transforms: Optional[list] = None, + images_paths:list[str]=None ) -> None: ''' Args: @@ -71,12 +71,19 @@ def __init__( tensor and expected target as input and returns a transformed version. These will be applied after albu_transforms. Defaults to None. + images_paths (list, optional): list of images to be considered. Overrides root_dir ''' super(FolderDataset, self).__init__(csv_file, root_dir, albu_transforms, end_transforms) - self.folder_images = [i for i in os.listdir(self.root_dir) - if i.endswith(('.JPG','.jpg','.JPEG','.jpeg'))] + if images_paths is None: + self.folder_images = [i for i in os.listdir(self.root_dir) + if i.endswith(('.JPG','.jpg','.JPEG','.jpeg'))] + else: + print("Overriding root_dir in FolderDataset, using images_paths") + self.folder_images = [i for i in images_paths + if i.endswith(('.JPG','.jpg','.JPEG','.jpeg'))] + self._img_names = self.folder_images self.anno_keys = self.data.columns diff --git a/animaloc/eval/evaluators.py b/animaloc/eval/evaluators.py index 89121f4..bebceae 100644 --- a/animaloc/eval/evaluators.py +++ b/animaloc/eval/evaluators.py @@ -176,6 +176,17 @@ def evaluate(self, returns: str = 'recall', wandb_flag: bool = False, viz: bool Returns: float ''' + + def batch_metrics(metric:Metrics,batchsize:int,output:dict)->None: + if batchsize>=1: + for i in range(batchsize): + gt = {k:v[i] for k,v in output['gt'].items()} + preds = {k:v[i] for k,v in output['preds'].items()} + counts = output['est_count'][i] + output_i = dict(gt = gt, preds = preds, est_count = counts) + metric.feed(**output_i) + else: + raise NotImplementedError self.model.eval() @@ -201,8 +212,12 @@ def evaluate(self, returns: str = 'recall', wandb_flag: bool = False, viz: bool wandb.log({'validation_vizuals': fig}) output = self.prepare_feeding(targets, output) + + batchsize = images.shape[0] - iter_metrics.feed(**output) + batch_metrics(iter_metrics,batchsize,output) + # iter_metrics.feed(**output) + iter_metrics.aggregate() if log_meters: logger.add_meter('n', sum(iter_metrics.tp) + sum(iter_metrics.fn)) @@ -226,7 +241,8 @@ def evaluate(self, returns: str = 'recall', wandb_flag: bool = False, viz: bool iter_metrics.flush() - self.metrics.feed(**output) + batch_metrics(self.metrics,batchsize,output) + # self.metrics.feed(**output) self._stored_metrics = self.metrics.copy() @@ -346,8 +362,12 @@ def post_stitcher(self, output: torch.Tensor) -> Any: def prepare_feeding(self, targets: Dict[str, torch.Tensor], output: List[torch.Tensor]) -> dict: - gt_coords = [p[::-1] for p in targets['points'].squeeze(0).tolist()] - gt_labels = targets['labels'].squeeze(0).tolist() + try: # batchsize==1 + gt_coords = [p[::-1] for p in targets['points'].cpu().tolist()] + gt_labels = targets['labels'].cpu().tolist() + except Exception: # batchsize>1 + gt_coords = [p[::-1] for p in targets['points']] + gt_labels = targets['labels'] gt = dict( loc = gt_coords, @@ -362,13 +382,14 @@ def prepare_feeding(self, targets: Dict[str, torch.Tensor], output: List[torch.T counts, locs, labels, scores, dscores = lmds(output) preds = dict( - loc = locs[0], - labels = labels[0], - scores = scores[0], - dscores = dscores[0] + loc = locs, + labels = labels, + scores = scores, + dscores = dscores ) - return dict(gt = gt, preds = preds, est_count = counts[0]) + + return dict(gt = gt, preds = preds, est_count = counts) @EVALUATORS.register() class DensityMapEvaluator(Evaluator): diff --git a/animaloc/models/dla.py b/animaloc/models/dla.py index 32b4457..2eb65b2 100644 --- a/animaloc/models/dla.py +++ b/animaloc/models/dla.py @@ -21,7 +21,6 @@ import math from os.path import join -from posixpath import basename import torch from torch import nn diff --git a/animaloc/models/faster_rcnn.py b/animaloc/models/faster_rcnn.py index ed6a847..800e954 100644 --- a/animaloc/models/faster_rcnn.py +++ b/animaloc/models/faster_rcnn.py @@ -90,7 +90,7 @@ def forward( if targets is not None: for t in targets: floating_point_types = (torch.float, torch.double, torch.half) - if not t["boxes"].dtype in floating_point_types: + if t["boxes"].dtype not in floating_point_types: raise TypeError(f"target boxes must of float type, instead got {t['boxes'].dtype}") if not t["labels"].dtype == torch.int64: raise TypeError("target labels must of int64 type, instead got {t['labels'].dtype}") diff --git a/animaloc/models/herdnet.py b/animaloc/models/herdnet.py index 87d437d..ad51f1c 100644 --- a/animaloc/models/herdnet.py +++ b/animaloc/models/herdnet.py @@ -19,7 +19,6 @@ import torch.nn as nn import numpy as np -import torchvision.transforms as T from typing import Optional diff --git a/animaloc/train/trainers.py b/animaloc/train/trainers.py index f6d09e2..dccfb9a 100644 --- a/animaloc/train/trainers.py +++ b/animaloc/train/trainers.py @@ -22,16 +22,13 @@ import wandb import matplotlib -import matplotlib.pyplot as plt matplotlib.use('Agg') -from torchvision.transforms import ToPILImage from typing import List, Optional, Union, Callable, Any from ..utils.torchvision_utils import SmoothedValue, reduce_dict from ..utils.logger import CustomLogger from ..eval.evaluators import Evaluator -from ..data.transforms import UnNormalize from .adaloss import Adaloss from ..utils.registry import Registry diff --git a/animaloc/vizual/objects.py b/animaloc/vizual/objects.py index df77a4b..6416f83 100644 --- a/animaloc/vizual/objects.py +++ b/animaloc/vizual/objects.py @@ -17,7 +17,7 @@ import PIL -from PIL import Image, ImageDraw, ImageFont +from PIL import ImageDraw, ImageFont __all__ = ['draw_points', 'draw_boxes', 'draw_text'] diff --git a/animaloc/vizual/plots.py b/animaloc/vizual/plots.py index 5560ae8..9994d5f 100644 --- a/animaloc/vizual/plots.py +++ b/animaloc/vizual/plots.py @@ -15,14 +15,11 @@ __version__ = "0.2.0" -import torch import matplotlib.pyplot as plt import random import itertools from typing import Optional -from torchvision.transforms import ToPILImage -from ..data.transforms import UnNormalize, GaussianMap __all__ = ['PlotPrecisionRecall'] diff --git a/my_ml_backend/Dockerfile b/my_ml_backend/Dockerfile new file mode 100644 index 0000000..d4ca437 --- /dev/null +++ b/my_ml_backend/Dockerfile @@ -0,0 +1,19 @@ +FROM python:3.8-slim + +RUN apt-get update && \ + apt-get upgrade -y && \ + apt-get install -y git + +ENV PYTHONUNBUFFERED=True \ + PORT=9090 \ + WORKERS=2 \ + THREADS=4 + +WORKDIR /app +COPY requirements.txt . + +RUN pip install --no-cache-dir -r requirements.txt + +COPY . ./ + +CMD exec gunicorn --preload --bind :$PORT --workers $WORKERS --threads $THREADS --timeout 0 _wsgi:app diff --git a/my_ml_backend/README.md b/my_ml_backend/README.md new file mode 100644 index 0000000..815a95d --- /dev/null +++ b/my_ml_backend/README.md @@ -0,0 +1,16 @@ +# Quickstart + +1. Build and start Machine Learning backend on `http://localhost:9090` + +```bash +docker-compose up +``` + +2. Validate that backend is running + +```bash +$ curl http://localhost:9090/health +{"status":"UP"} +``` + +3. Connect to the backend from Label Studio: go to your project `Settings -> Machine Learning -> Add Model` and specify `http://localhost:9090` as a URL. \ No newline at end of file diff --git a/my_ml_backend/_wsgi.py b/my_ml_backend/_wsgi.py new file mode 100644 index 0000000..9ab9ef5 --- /dev/null +++ b/my_ml_backend/_wsgi.py @@ -0,0 +1,114 @@ +import os +import argparse +import json +import logging +import logging.config + +logging.config.dictConfig({ + "version": 1, + "formatters": { + "standard": { + "format": "[%(asctime)s] [%(levelname)s] [%(name)s::%(funcName)s::%(lineno)d] %(message)s" + } + }, + "handlers": { + "console": { + "class": "logging.StreamHandler", + "level": os.getenv('LOG_LEVEL'), + "stream": "ext://sys.stdout", + "formatter": "standard" + } + }, + "root": { + "level": os.getenv('LOG_LEVEL'), + "handlers": [ + "console" + ], + "propagate": True + } +}) + +from label_studio_ml.api import init_app +from model import NewModel + + +_DEFAULT_CONFIG_PATH = os.path.join(os.path.dirname(__file__), 'config.json') + + +def get_kwargs_from_config(config_path=_DEFAULT_CONFIG_PATH): + if not os.path.exists(config_path): + return dict() + with open(config_path) as f: + config = json.load(f) + assert isinstance(config, dict) + return config + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description='Label studio') + parser.add_argument( + '-p', '--port', dest='port', type=int, default=9090, + help='Server port') + parser.add_argument( + '--host', dest='host', type=str, default='0.0.0.0', + help='Server host') + parser.add_argument( + '--kwargs', '--with', dest='kwargs', metavar='KEY=VAL', nargs='+', type=lambda kv: kv.split('='), + help='Additional LabelStudioMLBase model initialization kwargs') + parser.add_argument( + '-d', '--debug', dest='debug', action='store_true', + help='Switch debug mode') + parser.add_argument( + '--log-level', dest='log_level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'], default=None, + help='Logging level') + parser.add_argument( + '--model-dir', dest='model_dir', default=os.path.dirname(__file__), + help='Directory where models are stored (relative to the project directory)') + parser.add_argument( + '--check', dest='check', action='store_true', + help='Validate model instance before launching server') + + args = parser.parse_args() + + # setup logging level + if args.log_level: + logging.root.setLevel(args.log_level) + + def isfloat(value): + try: + float(value) + return True + except ValueError: + return False + + def parse_kwargs(): + param = dict() + for k, v in args.kwargs: + if v.isdigit(): + param[k] = int(v) + elif v == 'True' or v == 'true': + param[k] = True + elif v == 'False' or v == 'false': + param[k] = False + elif isfloat(v): + param[k] = float(v) + else: + param[k] = v + return param + + kwargs = get_kwargs_from_config() + + if args.kwargs: + kwargs.update(parse_kwargs()) + + if args.check: + print('Check "' + NewModel.__name__ + '" instance creation..') + model = NewModel(**kwargs) + + app = init_app(model_class=NewModel) + + app.run(host=args.host, port=args.port, debug=args.debug) + +else: + # for uWSGI use + app = init_app(model_class=NewModel) diff --git a/my_ml_backend/docker-compose.yml b/my_ml_backend/docker-compose.yml new file mode 100644 index 0000000..050e7f1 --- /dev/null +++ b/my_ml_backend/docker-compose.yml @@ -0,0 +1,13 @@ +version: "3.8" + +services: + ml-backend: + container_name: ml-backend + build: . + environment: + - MODEL_DIR=/data/models + - WORKERS=2 + - THREADS=4 + - LOG_LEVEL=DEBUG + ports: + - "9090:9090" diff --git a/my_ml_backend/model.py b/my_ml_backend/model.py new file mode 100644 index 0000000..1030fe8 --- /dev/null +++ b/my_ml_backend/model.py @@ -0,0 +1,158 @@ +from typing import List, Dict, Optional +from label_studio_ml.model import LabelStudioMLBase +from sahi.models.yolov8 import Yolov8DetectionModel +# from sahi.utils.cv import read_image +from sahi.predict import get_sliced_prediction, predict +from label_studio_ml.utils import (get_env, get_local_path, is_skipped) +from PIL import Image +import boto3 +import torch +from pathlib import Path +from urllib.parse import urlparse + +#labelstudio API settings +HOSTNAME = get_env('HOSTNAME', 'http://localhost:8080') +API_KEY = get_env("KEY") + +# Authenticate AWS +# PROFILE_NAME = 'my-profile' #TODO: update with your profile +# MY_SESSION = boto3.session.Session(profile_name=PROFILE_NAME) +# S3 = MY_SESSION.client('s3') + + +class Detector(object): + + def __init__(self, + path_to_weights:str, + confidence_threshold:float=0.3): + device = "cuda:0" if torch.cuda.is_available() else "cpu" + self.detection_model = Yolov8DetectionModel(model_path=path_to_weights, + confidence_threshold=confidence_threshold, + device=device) + print('Device:', device) + + def predict(self, urls:List): + + preds = list() + + for url in urls: + + # if data is on AWS + # r = urlparse(url, allow_fragments=False) + # bucket_name = r.netloc + # filename = r.path.lstrip('/') + # with open('./tmp/s3_img.jpg','wb+') as f: + # S3.download_fileobj(bucket_name, filename, f) + # img = Image.open(f) + + # if data is local + img = Image.open(get_local_path(url)) + + # get prediction + result = get_sliced_prediction(img, + self.detection_model, + slice_height=640, + slice_width=640, + overlap_height_ratio=0.1, + overlap_width_ratio=0.1, + postprocess_type='NMS', + ) + img_height = result.image_height + img_width = result.image_width + formatted_pred = [self.format_prediction(pred, + img_height=img_height, + img_width=img_width) for pred in result.to_coco_annotations()] + preds.append({'result':formatted_pred}) + + return preds + + def format_prediction(self,pred:Dict,img_height:int,img_width:int): + # formatting the prediction to work with Label studio + x, y, width, height = pred['bbox'] + label = pred['category_name'] + score = pred['score'] + + template = { + "from_name": "label", + "to_name": "image", + "type": "rectanglelabels", + 'value': { + 'rectanglelabels': [label], + 'x': x / img_width * 100, + 'y': y / img_height * 100, + 'width': width / img_width * 100, + 'height': height / img_height * 100 + }, + 'score': score + } + + return template + + def train(self, dataloader): + pass + +class NewModel(LabelStudioMLBase): + + def __init__(self,project_id,label_config,**kwargs): + super(NewModel, self).__init__(project_id=project_id, + label_config=label_config, + **kwargs) + self.from_name = "label" + self.to_name = "image" + self.value = "image" + + # Load localizer + self.model = Detector(path_to_weights=r"C:\Users\Machine Learning\Desktop\workspace-wildAI\yolov8.kaza.pt", + confidence_threshold=0.3) + + def predict(self, tasks: List[Dict], context: Optional[Dict] = None, **kwargs) -> List[Dict]: + """ Write your inference logic here + :param tasks: [Label Studio tasks in JSON format](https://labelstud.io/guide/task_format.html) + :param context: [Label Studio context in JSON format](https://labelstud.io/guide/ml.html#Passing-data-to-ML-backend) + :return predictions: [Predictions array in JSON format](https://labelstud.io/guide/export.html#Raw-JSON-format-of-completed-tasks) + """ + # print(f'''\ + # * Run prediction on {tasks} + + # * Received context: {context} + + # * Project ID: {self.project_id} + + # * Label config: {self.label_config} + + # * Parsed JSON Label config: {self.parsed_label_config}''') + + + image_urls = [task['data'][self.value] for task in tasks] + predictions = self.model.predict(image_urls) + + # print(f"Predictions: {predictions}\n") + + return predictions + + def fit(self, event, data, **kwargs): + """ + This method is called each time an annotation is created or updated + You can run your logic here to update the model and persist it to the cache + It is not recommended to perform long-running operations here, as it will block the main thread + Instead, consider running a separate process or a thread (like RQ worker) to perform the training + :param event: event type can be ('ANNOTATION_CREATED', 'ANNOTATION_UPDATED') + :param data: the payload received from the event (check [Webhook event reference](https://labelstud.io/guide/webhook_reference.html)) + """ + + # use cache to retrieve the data from the previous fit() runs + # old_data = self.get('my_data') + # old_model_version = self.get('model_version') + # print(f'Old data: {old_data}') + # print(f'Old model version: {old_model_version}') + + # store new data to the cache + # self.set('my_data', 'my_new_data_value') + # self.set('model_version', 'my_new_model_version') + # print(f'New data: {self.get("my_data")}') + # print(f'New model version: {self.get("model_version")}') + + # print('fit() completed successfully.') + + pass + diff --git a/my_ml_backend/requirements.txt b/my_ml_backend/requirements.txt new file mode 100644 index 0000000..5bf0cd8 --- /dev/null +++ b/my_ml_backend/requirements.txt @@ -0,0 +1,6 @@ +gunicorn==20.1.0 +label-studio-ml @ git+https://github.com/HumanSignal/label-studio-ml-backend.git +ultralytics +torch -index-url https://download.pytorch.org/whl/cu118 +sahi +boto3 \ No newline at end of file diff --git a/notebooks/demo-training-testing-herdnet.ipynb b/notebooks/demo-training-testing-herdnet.ipynb index 61b9481..dbbd8bf 100644 --- a/notebooks/demo-training-testing-herdnet.ipynb +++ b/notebooks/demo-training-testing-herdnet.ipynb @@ -1 +1,480 @@ -{"cells":[{"cell_type":"markdown","metadata":{"id":"tOMmy2YvTHSy"},"source":["# DEMO - Training and testing HerdNet on nadir aerial images"]},{"cell_type":"markdown","metadata":{"id":"_XTpIbRwT9PO"},"source":["## Installations"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"9v5ab5BbSrVl"},"outputs":[],"source":["# Check GPU\n","!nvidia-smi"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"z9RVoQx5UOkg"},"outputs":[],"source":["# Install the dependencies\n","!pip install albumentations==1.0.3\n","!pip install fiftyone==0.14.3\n","!pip install hydra-core==1.1.0\n","!pip install opencv-python==4.5.1.48\n","!pip install pandas==1.2.3\n","!pip install pillow==8.2.0\n","!pip install scikit-image==0.18.1\n","!pip install scikit-learn==1.0.2\n","!pip install scipy==1.6.2\n","!pip install wandb==0.10.33"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"oJydpjxoUGAC"},"outputs":[],"source":["# Download and install the code\n","import sys\n","\n","!git clone https://github.com/Alexandre-Delplanque/HerdNet\n","!cd '/content/HerdNet' && python setup.py install\n","\n","sys.path.append('/content/HerdNet')"]},{"cell_type":"markdown","metadata":{"id":"cuxaC9qGVk5S"},"source":["## Create datasets"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"0yt06VTPVtsq"},"outputs":[],"source":["# Download some of the data of Delplanque et al. (2021) as an example\n","!gdown 1CcTAZZJdwrBfCPJtVH6VBU3luGKIN9st -O /content/data.zip\n","!unzip -oq /content/data.zip -d /content"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"y1n_cQLFYBsJ"},"outputs":[],"source":["# Set the seed\n","from animaloc.utils.seed import set_seed\n","\n","set_seed(9292)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"9tGqtLNG21Jf"},"outputs":[],"source":["# Create validation patches using the patcher tool (for demo)\n","from animaloc.utils.useful_funcs import mkdir\n","\n","mkdir('/content/data/val_patches')\n","!python /content/HerdNet/tools/patcher.py /content/data/val 512 512 0 /content/data/val_patches -csv /content/data/val.csv -min 0.0 -all False"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"Bwp4XPR8YNMR"},"outputs":[],"source":["# Training, validation and test datasets\n","import albumentations as A\n","\n","from animaloc.datasets import CSVDataset\n","from animaloc.data.transforms import MultiTransformsWrapper, DownSample, PointsToMask, FIDT\n","\n","patch_size = 512\n","num_classes = 7\n","down_ratio = 2\n","\n","train_dataset = CSVDataset(\n"," csv_file = '/content/data/train_patches.csv',\n"," root_dir = '/content/data/train_patches',\n"," albu_transforms = [\n"," A.VerticalFlip(p=0.5), \n"," A.HorizontalFlip(p=0.5),\n"," A.RandomRotate90(p=0.5),\n"," A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.2),\n"," A.Blur(blur_limit=15, p=0.2),\n"," A.Normalize(p=1.0)\n"," ],\n"," end_transforms = [MultiTransformsWrapper([\n"," FIDT(num_classes=num_classes, down_ratio=down_ratio),\n"," PointsToMask(radius=2, num_classes=num_classes, squeeze=True, down_ratio=int(patch_size//16))\n"," ])]\n"," )\n","\n","val_dataset = CSVDataset(\n"," csv_file = '/content/data/val_patches/gt.csv',\n"," root_dir = '/content/data/val_patches',\n"," albu_transforms = [A.Normalize(p=1.0)],\n"," end_transforms = [DownSample(down_ratio=down_ratio, anno_type='point')]\n"," )\n","\n","test_dataset = CSVDataset(\n"," csv_file = '/content/data/test.csv',\n"," root_dir = '/content/data/test',\n"," albu_transforms = [A.Normalize(p=1.0)],\n"," end_transforms = [DownSample(down_ratio=down_ratio, anno_type='point')]\n"," )"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"lR1W5NVFYhiZ"},"outputs":[],"source":["# Dataloaders\n","from torch.utils.data import DataLoader\n","\n","train_dataloader = DataLoader(dataset = train_dataset, batch_size = 4, shuffle = True)\n","\n","val_dataloader = DataLoader(dataset = val_dataset, batch_size = 1, shuffle = False)\n","\n","test_dataloader = DataLoader(dataset = test_dataset, batch_size = 1, shuffle = False)"]},{"cell_type":"markdown","metadata":{"id":"emWQUMq2Vwpj"},"source":["## Define HerdNet for training"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"JIBKygFlV0V1"},"outputs":[],"source":["from animaloc.models import HerdNet\n","from torch import Tensor\n","from animaloc.models import LossWrapper\n","from animaloc.train.losses import FocalLoss\n","from torch.nn import CrossEntropyLoss\n","\n","herdnet = HerdNet(num_classes=num_classes, down_ratio=down_ratio).cuda()\n","\n","weight = Tensor([0.1, 1.0, 2.0, 1.0, 6.0, 12.0, 1.0]).cuda()\n","\n","losses = [\n"," {'loss': FocalLoss(reduction='mean'), 'idx': 0, 'idy': 0, 'lambda': 1.0, 'name': 'focal_loss'},\n"," {'loss': CrossEntropyLoss(reduction='mean', weight=weight), 'idx': 1, 'idy': 1, 'lambda': 1.0, 'name': 'ce_loss'}\n"," ]\n","\n","herdnet = LossWrapper(herdnet, losses=losses)"]},{"cell_type":"markdown","metadata":{"id":"Nm5u6yg4V78C"},"source":["## Create the Trainer"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"MSBimwtzWDZp"},"outputs":[],"source":["from torch.optim import Adam\n","from animaloc.train import Trainer\n","from animaloc.eval import PointsMetrics, HerdNetStitcher, HerdNetEvaluator\n","from animaloc.utils.useful_funcs import mkdir\n","\n","work_dir = '/content/output'\n","mkdir(work_dir)\n","\n","lr = 1e-4\n","weight_decay = 1e-3\n","epochs = 100\n","\n","optimizer = Adam(params=herdnet.parameters(), lr=lr, weight_decay=weight_decay)\n","\n","metrics = PointsMetrics(radius=20, num_classes=num_classes)\n","\n","stitcher = HerdNetStitcher(\n"," model=herdnet, \n"," size=(patch_size,patch_size), \n"," overlap=160, \n"," down_ratio=down_ratio, \n"," reduction='mean'\n"," )\n","\n","evaluator = HerdNetEvaluator(\n"," model=herdnet, \n"," dataloader=val_dataloader, \n"," metrics=metrics, \n"," stitcher=stitcher, \n"," work_dir=work_dir, \n"," header='validation'\n"," )\n","\n","trainer = Trainer(\n"," model=herdnet,\n"," train_dataloader=train_dataloader,\n"," optimizer=optimizer,\n"," num_epochs=epochs,\n"," evaluator=evaluator,\n"," work_dir=work_dir\n"," )"]},{"cell_type":"markdown","metadata":{"id":"axsTtq4WV0ot"},"source":["## Start training"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"malFT6r5V4rC"},"outputs":[],"source":["trainer.start(warmup_iters=100, checkpoints='best', select='max', validate_on='f1_score')"]},{"cell_type":"markdown","metadata":{"id":"_e0CQd5Bxx5T"},"source":["## Test the model"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"D5133GoRz8r_"},"outputs":[],"source":["# Path to your .pth file\n","import gdown\n","\n","pth_path = ''\n","\n","if not pth_path:\n"," gdown.download(\n"," 'https://drive.google.com/uc?export=download&id=1-WUnBC4BJMVkNvRqalF_HzA1_pRkQTI_',\n"," '/content/20220413_herdnet_model.pth'\n"," )\n"," pth_path = '/content/20220413_herdnet_model.pth'"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"VPHXVYWNzVDj"},"outputs":[],"source":["# Create output folder\n","test_dir = '/content/test_output'\n","mkdir(test_dir)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"DXUsHk7dzl47"},"outputs":[],"source":["# Load trained parameters\n","from animaloc.models import load_model\n","\n","herdnet = load_model(herdnet, pth_path=pth_path)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"lX3Jp883zB-D"},"outputs":[],"source":["# Create an Evaluator\n","test_evaluator = HerdNetEvaluator(\n"," model=herdnet, \n"," dataloader=test_dataloader, \n"," metrics=metrics, \n"," stitcher=stitcher, \n"," work_dir=test_dir, \n"," header='test'\n"," )"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"CZt_wNle0488"},"outputs":[],"source":["# Start testing\n","test_f1_score = test_evaluator.evaluate(returns='f1_score')"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"1vICHF-sFGxa"},"outputs":[],"source":["# Print global F1 score (%)\n","print(f\"F1 score = {test_f1_score * 100:0.0f}%\")"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"ANdn_feR2ZY8"},"outputs":[],"source":["# Get the detections\n","detections = test_evaluator.results\n","detections"]}],"metadata":{"accelerator":"GPU","colab":{"authorship_tag":"ABX9TyOEeDjFeLKZlLHdeq13OrD5","provenance":[]},"gpuClass":"standard","kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"name":"python"}},"nbformat":4,"nbformat_minor":0} +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "tOMmy2YvTHSy" + }, + "source": [ + "# DEMO - Training and testing HerdNet on nadir aerial images" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_XTpIbRwT9PO" + }, + "source": [ + "## Installations" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9v5ab5BbSrVl" + }, + "outputs": [], + "source": [ + "# Check GPU\n", + "!nvidia-smi" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "z9RVoQx5UOkg" + }, + "outputs": [], + "source": [ + "# Install the dependencies\n", + "!pip install albumentations==1.0.3\n", + "!pip install fiftyone==0.14.3\n", + "!pip install hydra-core==1.1.0\n", + "!pip install opencv-python==4.5.1.48\n", + "!pip install pandas==1.2.3\n", + "!pip install pillow==8.2.0\n", + "!pip install scikit-image==0.18.1\n", + "!pip install scikit-learn==1.0.2\n", + "!pip install scipy==1.6.2\n", + "!pip install wandb==0.10.33" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "oJydpjxoUGAC" + }, + "outputs": [], + "source": [ + "# Download and install the code\n", + "import sys\n", + "\n", + "!git clone https://github.com/Alexandre-Delplanque/HerdNet\n", + "!cd '/content/HerdNet' && python setup.py install\n", + "\n", + "sys.path.append('/content/HerdNet')" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "cuxaC9qGVk5S" + }, + "source": [ + "## Create datasets" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "0yt06VTPVtsq" + }, + "outputs": [], + "source": [ + "# Download some of the data of Delplanque et al. (2021) as an example\n", + "!gdown 1CcTAZZJdwrBfCPJtVH6VBU3luGKIN9st -O /content/data.zip\n", + "!unzip -oq /content/data.zip -d /content" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "y1n_cQLFYBsJ" + }, + "outputs": [], + "source": [ + "# Set the seed\n", + "from animaloc.utils.seed import set_seed\n", + "\n", + "set_seed(9292)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9tGqtLNG21Jf" + }, + "outputs": [], + "source": [ + "# Create validation patches using the patcher tool (for demo)\n", + "from animaloc.utils.useful_funcs import mkdir\n", + "\n", + "mkdir('/content/data/val_patches')\n", + "!python /content/HerdNet/tools/patcher.py /content/data/val 512 512 0 /content/data/val_patches -csv /content/data/val.csv -min 0.0 -all False" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Bwp4XPR8YNMR" + }, + "outputs": [], + "source": [ + "# Training, validation and test datasets\n", + "import albumentations as A\n", + "\n", + "from animaloc.datasets import CSVDataset, FolderDataset\n", + "from animaloc.data.transforms import MultiTransformsWrapper, DownSample, PointsToMask, FIDT\n", + "\n", + "patch_size = 512\n", + "num_classes = 7\n", + "down_ratio = 2\n", + "# FolderDataset allows to account for empty images if present.\n", + "train_dataset = FolderDataset(\n", + " csv_file = '/content/data/train_patches.csv',\n", + " root_dir = '/content/data/train_patches',\n", + " albu_transforms = [\n", + " A.VerticalFlip(p=0.5), \n", + " A.HorizontalFlip(p=0.5),\n", + " A.RandomRotate90(p=0.5),\n", + " A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.2),\n", + " A.Blur(blur_limit=15, p=0.2),\n", + " A.Normalize(p=1.0)\n", + " ],\n", + " end_transforms = [MultiTransformsWrapper([\n", + " FIDT(num_classes=num_classes, down_ratio=down_ratio),\n", + " PointsToMask(radius=2, num_classes=num_classes, squeeze=True, down_ratio=int(patch_size//16))\n", + " ])]\n", + " )\n", + "\n", + "val_dataset = FolderDataset(\n", + " csv_file = '/content/data/val_patches/gt.csv',\n", + " root_dir = '/content/data/val_patches',\n", + " albu_transforms = [A.Normalize(p=1.0)],\n", + " end_transforms = [DownSample(down_ratio=down_ratio, anno_type='point')]\n", + " )\n", + "\n", + "test_dataset = FolderDataset(\n", + " csv_file = '/content/data/test.csv',\n", + " root_dir = '/content/data/test',\n", + " albu_transforms = [A.Normalize(p=1.0)],\n", + " end_transforms = [DownSample(down_ratio=down_ratio, anno_type='point')]\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def val_collate_fn(batch: tuple)->tuple[torch.Tensor,dict]:\n", + " \"\"\"collate_fn used to create the validation dataloader\n", + "\n", + " Args:\n", + " batch (tuple): (img:torch.Tensor, targets:dict)\n", + "\n", + " Returns:\n", + " tuple: (image, target)\n", + " \"\"\"\n", + "\n", + " batched = dict(points=[], labels=[])\n", + " batch_img = torch.stack([p[0] for p in batch])\n", + " targets = [p[1] for p in batch]\n", + " keys = targets[0].keys()\n", + "\n", + " # get non_empty samples indidces\n", + " non_empty_idx = [i for i, a in enumerate(targets) if len(a[\"labels\"]) > 0]\n", + " targets_empty = [\n", + " targets[i] for i in list(set(range(len(batch))) - set(non_empty_idx))\n", + " ]\n", + " targets = [targets[i] for i in non_empty_idx]\n", + "\n", + " # Creating batch\n", + " for k in keys:\n", + " batched[k] = [] # initialize to be empty list\n", + " if k == \"points\":\n", + " batched[k] = [a[k].cpu().tolist() for a in targets]\n", + " if len(targets_empty) > 0:\n", + " batched[k] = batched[k] + [[]] * len(targets_empty)\n", + " if k == \"labels\":\n", + " batched[k] = [a[k].cpu().tolist() for a in targets]\n", + " # batched[k] = [a if isinstance(a,list) else [a] for a in batched[k]]\n", + " if len(targets_empty) > 0:\n", + " batched[k] = batched[k] + [[]] * len(targets_empty)\n", + "\n", + " return batch_img, batched" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "lR1W5NVFYhiZ" + }, + "outputs": [], + "source": [ + "# Dataloaders\n", + "from torch.utils.data import DataLoader\n", + "\n", + "train_dataloader = DataLoader(dataset = train_dataset, batch_size = 4, shuffle = True)\n", + "\n", + "val_dataloader = DataLoader(dataset = val_dataset, batch_size = 1, shuffle = False, collate_fn=val_collate_fn)\n", + "\n", + "test_dataloader = DataLoader(dataset = test_dataset, batch_size = 1, shuffle = False, collate_fn=val_collate_fn)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "emWQUMq2Vwpj" + }, + "source": [ + "## Define HerdNet for training" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "JIBKygFlV0V1" + }, + "outputs": [], + "source": [ + "from animaloc.models import HerdNet\n", + "from torch import Tensor\n", + "from animaloc.models import LossWrapper\n", + "from animaloc.train.losses import FocalLoss\n", + "from torch.nn import CrossEntropyLoss\n", + "\n", + "herdnet = HerdNet(num_classes=num_classes, down_ratio=down_ratio).cuda()\n", + "\n", + "weight = Tensor([0.1, 1.0, 2.0, 1.0, 6.0, 12.0, 1.0]).cuda()\n", + "\n", + "losses = [\n", + " {'loss': FocalLoss(reduction='mean'), 'idx': 0, 'idy': 0, 'lambda': 1.0, 'name': 'focal_loss'},\n", + " {'loss': CrossEntropyLoss(reduction='mean', weight=weight), 'idx': 1, 'idy': 1, 'lambda': 1.0, 'name': 'ce_loss'}\n", + " ]\n", + "\n", + "herdnet = LossWrapper(herdnet, losses=losses)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Nm5u6yg4V78C" + }, + "source": [ + "## Create the Trainer" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "MSBimwtzWDZp" + }, + "outputs": [], + "source": [ + "from torch.optim import Adam\n", + "from animaloc.train import Trainer\n", + "from animaloc.eval import PointsMetrics, HerdNetStitcher, HerdNetEvaluator\n", + "from animaloc.utils.useful_funcs import mkdir\n", + "\n", + "work_dir = '/content/output'\n", + "mkdir(work_dir)\n", + "\n", + "lr = 1e-4\n", + "weight_decay = 1e-3\n", + "epochs = 100\n", + "\n", + "optimizer = Adam(params=herdnet.parameters(), lr=lr, weight_decay=weight_decay)\n", + "\n", + "metrics = PointsMetrics(radius=20, num_classes=num_classes)\n", + "\n", + "stitcher = HerdNetStitcher(\n", + " model=herdnet, \n", + " size=(patch_size,patch_size), \n", + " overlap=160, \n", + " down_ratio=down_ratio, \n", + " reduction='mean'\n", + " )\n", + "\n", + "evaluator = HerdNetEvaluator(\n", + " model=herdnet, \n", + " dataloader=val_dataloader, \n", + " metrics=metrics, \n", + " stitcher=stitcher, \n", + " work_dir=work_dir, \n", + " header='validation'\n", + " )\n", + "\n", + "trainer = Trainer(\n", + " model=herdnet,\n", + " train_dataloader=train_dataloader,\n", + " optimizer=optimizer,\n", + " num_epochs=epochs,\n", + " evaluator=evaluator,\n", + " work_dir=work_dir\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "axsTtq4WV0ot" + }, + "source": [ + "## Start training" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "malFT6r5V4rC" + }, + "outputs": [], + "source": [ + "trainer.start(warmup_iters=100, checkpoints='best', select='max', validate_on='f1_score')" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_e0CQd5Bxx5T" + }, + "source": [ + "## Test the model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "D5133GoRz8r_" + }, + "outputs": [], + "source": [ + "# Path to your .pth file\n", + "import gdown\n", + "\n", + "pth_path = ''\n", + "\n", + "if not pth_path:\n", + " gdown.download(\n", + " 'https://drive.google.com/uc?export=download&id=1-WUnBC4BJMVkNvRqalF_HzA1_pRkQTI_',\n", + " '/content/20220413_herdnet_model.pth'\n", + " )\n", + " pth_path = '/content/20220413_herdnet_model.pth'" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "VPHXVYWNzVDj" + }, + "outputs": [], + "source": [ + "# Create output folder\n", + "test_dir = '/content/test_output'\n", + "mkdir(test_dir)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "DXUsHk7dzl47" + }, + "outputs": [], + "source": [ + "# Load trained parameters\n", + "from animaloc.models import load_model\n", + "\n", + "herdnet = load_model(herdnet, pth_path=pth_path)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "lX3Jp883zB-D" + }, + "outputs": [], + "source": [ + "# Create an Evaluator\n", + "test_evaluator = HerdNetEvaluator(\n", + " model=herdnet, \n", + " dataloader=test_dataloader, \n", + " metrics=metrics, \n", + " stitcher=stitcher, \n", + " work_dir=test_dir, \n", + " header='test'\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "CZt_wNle0488" + }, + "outputs": [], + "source": [ + "# Start testing\n", + "test_f1_score = test_evaluator.evaluate(returns='f1_score')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "1vICHF-sFGxa" + }, + "outputs": [], + "source": [ + "# Print global F1 score (%)\n", + "print(f\"F1 score = {test_f1_score * 100:0.0f}%\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ANdn_feR2ZY8" + }, + "outputs": [], + "source": [ + "# Get the detections\n", + "detections = test_evaluator.results\n", + "detections" + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "authorship_tag": "ABX9TyOEeDjFeLKZlLHdeq13OrD5", + "provenance": [] + }, + "gpuClass": "standard", + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..65b216a --- /dev/null +++ b/requirements.txt @@ -0,0 +1,19 @@ +albumentations +fiftyone +hydra-core +opencv-python +pandas +pillow +scikit-image +scikit-learn +scipy +wandb +torch +torchvision +boto3 +sagemaker +rasterio +tqdm +ultralytics +geopandas +sahi \ No newline at end of file diff --git a/tools/patcher.py b/tools/patcher.py index 9226c78..f1290c2 100644 --- a/tools/patcher.py +++ b/tools/patcher.py @@ -21,8 +21,12 @@ import torchvision import numpy import cv2 +from pathlib import Path +import math from albumentations import PadIfNeeded +from itertools import product +from torchvision.utils import save_image from tqdm import tqdm @@ -33,12 +37,20 @@ parser.add_argument('root', type=str, help='path to the images directory (str)') parser.add_argument('height', type=int, - help='height of the patches, in pixels (int)') + help='height of the patches, in pixels (int)' + ) parser.add_argument('width', type=int, - help='width of the patches, in pixels (int)') + help='width of the patches, in pixels (int)' + ) parser.add_argument('overlap', type=int, help='overlap between patches, in pixels (int)') -parser.add_argument('dest', type=str, +parser.add_argument('-overlapfactor', type=float,default=0.0, + help='overlap ratio between patches, in [0,1]. It works only for raw images. It does not work with -csv') +parser.add_argument('-ratioheight',type=float,default=0.0, + help='ratios for height. When it is not 1.0 then the height of the tile is infered from the image height and the ratio.') +parser.add_argument('-ratiowidth',type=float,default=0.0, + help='ratios for width. When it is not 1.0 then the width of the tile is infered from the image width and the ratio.') +parser.add_argument('-dest', type=str, help='destination path (str)') parser.add_argument('-csv', type=str, help='path to a csv file containing annotations (str). Defaults to None') @@ -46,27 +58,132 @@ help='minimum fraction of area for an annotation to be kept (float). Defautls to 0.1') parser.add_argument('-all', type=bool, default=False, help='set to True to save all patches, not only those containing annotations (bool). Defaults to False') +parser.add_argument('-pattern',type=str,default='**/*.jpg', + help='pattern of files extension') +parser.add_argument('-rmheight',type=float,default=0.0, + help='height overlap to be removed at both bottom and top') +parser.add_argument('-rmwidth',type=float,default=0.0, + help='width overlap to be removed at both left and right sides') + args = parser.parse_args() +#Helper funcs +def get_patches(image,tile_w:int,tile_h:int,overlaping_factor:float): + + patches = list() + image_width=image.shape[2] + image_height=image.shape[1] + + def get_coordinates(): + + # x limits + lim = math.ceil((image_width-tile_w)/((1-overlaping_factor)*tile_w)) + x_right = [math.floor(tile_w + i*(1-overlaping_factor)*tile_w) for i in range(lim)] + x_coords = [(x-tile_w,x) for x in x_right] + if len(x_coords)>0: + left,right = x_coords[-1] + x_coords[-1] = (left,image_width) # extending to remaining pixels + + # y limits + lim = math.ceil((image_height-tile_h)/((1-overlaping_factor)*tile_h)) + y_bottom = [math.floor(tile_h + i*(1-overlaping_factor)*tile_h) for i in range(lim)] + y_coords = [(y-tile_h,y) for y in y_bottom] + if len(y_coords)>0: + top,bottom = y_coords[-1] + y_coords[-1] = (top,image_height) # extending to remaining pixels + + # tiles coordinates + if len(y_coords)>0 and len(x_coords)>0: + pass + elif len(y_coords) == 0: + y_coords = [(0,image_height),] + elif len(x_coords) == 0: + x_coords = [(0,image_width),] + + coordinates = product(x_coords,y_coords) + return list(coordinates) + + coords = get_coordinates() + + # store patches + for (x_left,x_right),(y_top,y_bottom) in coords: + patches.append(image[:,y_top:y_bottom,x_left:x_right]) + + return patches + +def save_list_images( + batch:list, + basename: str, + dest_folder: str + ) -> None: + ''' Save mini-batch tensors into image files + + Use torchvision save_image function, + see https://pytorch.org/vision/stable/utils.html#torchvision.utils.save_image + + Args: + batch (list): mini-batch tensor + basename (str) : parent image name, with extension + dest_folder (str): destination folder path + ''' + + base_wo_extension, extension = basename.split('.')[0], basename.split('.')[1] + for i, b in enumerate(range(len(batch))): + full_path = '_'.join([base_wo_extension, str(i) + '.']) + extension + save_path = os.path.join(dest_folder, full_path) + save_image(batch[b], fp=save_path) + + def main(): - images_paths = [os.path.join(args.root, p) for p in os.listdir(args.root) if not p.endswith('.csv')] + # images_paths =list(Path(args.root).glob(args.pattern)) #[os.path.join(args.root, p) for p in os.listdir(args.root) if not p.endswith('.csv')] + images_paths = [p for p in Path(args.root).glob(args.pattern)] if args.csv is not None: patches_buffer = PatchesBuffer(args.csv, args.root, (args.height, args.width), overlap=args.overlap, min_visibility=args.min).buffer patches_buffer.drop(columns='limits').to_csv(os.path.join(args.dest, 'gt.csv'), index=False) for img_path in tqdm(images_paths, desc='Exporting patches'): - pil_img = PIL.Image.open(img_path) + try: + pil_img = PIL.Image.open(img_path) + except : + print("failed for: ",img_path,flush=True) + continue img_tensor = torchvision.transforms.ToTensor()(pil_img) img_name = os.path.basename(img_path) + + # Cropping out image-level overlap + height_overlap = math.ceil(args.rmheight * img_tensor.shape[1]) + width_overlap = math.ceil(args.rmwidth * img_tensor.shape[2]) + if height_overlap*width_overlap > 0 : + img_tensor = img_tensor[:,height_overlap:-height_overlap, width_overlap:-width_overlap] + print(f"Removing {2*width_overlap} pixels to the width; and {2*height_overlap} pixels to the height.") + elif (height_overlap == 0) and (width_overlap != 0): + img_tensor = img_tensor[:,:, width_overlap:-width_overlap] + elif width_overlap == 0 and (height_overlap != 0): + img_tensor = img_tensor[:,height_overlap:-height_overlap,:] + + # Computes tile width and height using the given ratios + # It overrides the parameters 'width' and 'height' + assert (args.ratiowidth<=1.0) and (args.ratioheight<=1.0), "The ratios should be at most 1.0" + if args.ratiowidth > 0.0: + args.width = math.ceil(img_tensor.shape[2]*args.ratiowidth) + if args.ratioheight > 0.0: + args.height = math.ceil(img_tensor.shape[1]*args.ratioheight) + # checking overlapfactor provided + assert args.overlapfactor<1, 'It should be less than 1.' + if args.csv is not None: # save all patches if args.all: - patches = ImageToPatches(img_tensor, (args.height, args.width), overlap=args.overlap).make_patches() - save_batch_images(patches, img_name, args.dest) + if args.overlapfactor>0 : + patches = get_patches(img_tensor,tile_w=args.width,tile_h=args.height,overlaping_factor=args.overlapfactor) + save_batch_images(patches, img_name, args.dest) + else: + patches = ImageToPatches(img_tensor, (args.height, args.width), overlap=args.overlap).make_patches() + save_list_images(patches, img_name, args.dest) # or only annotated ones else: @@ -84,8 +201,12 @@ def main(): padded_img.save(os.path.join(args.dest, ptch_name)) else: - patches = ImageToPatches(img_tensor, (args.height, args.width), overlap=args.overlap).make_patches() - save_batch_images(patches, img_name, args.dest) + if args.overlapfactor > 0: + patches = get_patches(img_tensor,tile_w=args.width,tile_h=args.height,overlaping_factor=args.overlapfactor) + save_list_images(patches, img_name, args.dest) + else: + patches = ImageToPatches(img_tensor, (args.height, args.width), overlap=args.overlap).make_patches() + save_batch_images(patches, img_name, args.dest) if __name__ == '__main__':