diff --git a/directai_fastapi/Dockerfile b/directai_fastapi/Dockerfile index d495a33..d18603f 100644 --- a/directai_fastapi/Dockerfile +++ b/directai_fastapi/Dockerfile @@ -15,6 +15,8 @@ RUN apt-get install libjpeg-dev zlib1g-dev -y RUN pip uninstall -y pillow RUN CC="cc -mavx2" pip install -U --force-reinstall pillow-simd +RUN apt-get install wget unzip -y + COPY logging_config.py . COPY pydantic_models.py . COPY server.py . @@ -29,6 +31,9 @@ COPY unit_tests unit_tests RUN mkdir modeling COPY modeling modeling +RUN mkdir benchmarking +COPY benchmarking benchmarking + ENV PYTHONPATH "${PYTHONPATH}:/directai_fastapi" RUN chmod +x run.sh diff --git a/directai_fastapi/benchmarking/benchmark_coco.py b/directai_fastapi/benchmarking/benchmark_coco.py new file mode 100644 index 0000000..b7459c6 --- /dev/null +++ b/directai_fastapi/benchmarking/benchmark_coco.py @@ -0,0 +1,121 @@ +import os +from pycocotools.coco import COCO +from pycocotools.cocoeval import COCOeval +import json + +from logging_config import logger +from modeling.batch_processing import ( + build_ray_dataset_from_directory, + run_object_detector_against_ray_dataset, + convert_detections_to_coco_format, + build_naive_detector_config_from_coco_categories_metadata, + filter_dataset_by_path, +) + + +def download_coco_data( + split: str, to_dir: str = "/directai_fastapi/.cache/coco" +) -> None: + """Download COCO data for the given split if not present.""" + # Define your paths and URLs for COCO datasets + if split not in ["train", "val", "test"]: + raise ValueError(f"Invalid split: {split}. Expected 'train', 'val', or 'test'.") + + annotation_url = ( + f"http://images.cocodataset.org/annotations/annotations_trainval2017.zip" + ) + images_url = f"http://images.cocodataset.org/zips/{split}2017.zip" + + annotation_zip_file = os.path.join(to_dir, f"annotations_trainval2017.zip") + images_zip_file = os.path.join(to_dir, f"{split}2017.zip") + + # Check if files exist, download if not + if not os.path.exists(annotation_zip_file) or not os.path.exists(images_zip_file): + os.makedirs(to_dir, exist_ok=True) + # Download annotations + os.system(f"wget {annotation_url} -O {annotation_zip_file}") + os.system(f"unzip -o {annotation_zip_file} -d {to_dir}") + # Download images + os.system(f"wget {images_url} -O {images_zip_file}") + os.system(f"unzip -o {images_zip_file} -d {to_dir}") + + +if __name__ == "__main__": + coco_base_dir = "/directai_fastapi/.cache/coco" + split = "val" + + logger.info(f"Downloading COCO data for split: {split}") + download_coco_data(split, coco_base_dir) + + images_dir = os.path.join(coco_base_dir, f"{split}2017") + annotations_path = os.path.join( + coco_base_dir, "annotations", f"instances_{split}2017.json" + ) + predictions_path = os.path.join(coco_base_dir, f"predictions_{split}2017.json") + + # Download COCO data + download_coco_data(split, to_dir=coco_base_dir) + + # Load COCO ground truth + coco_gt = COCO(annotations_path) + + img_ids = coco_gt.getImgIds() + # img_ids = img_ids[:10] # Limit to 100 images for testing + + # for the COCO gt object, the cats and imgs attributes are dictionaries with id as key + # instead of the list of dictionaries that is the supposed format + # so we convert them to a list of dictionaries + categories_metadata = [c for c in coco_gt.cats.values()] + images_metadata = [i for i in coco_gt.imgs.values()] + + paths_set = set([coco_gt.imgs[id]["file_name"] for id in img_ids]) + + # Build naive detector config from COCO categories metadata + labels, inc_sub_labels_dict = ( + build_naive_detector_config_from_coco_categories_metadata(categories_metadata) + ) + label_conf_thres = {name: 0.001 for name in labels} + + # print(labels) + # print(categories_metadata) + # print(coco_gt.cats) + # print(img_ids) + + # Run object detector against COCO dataset + ray_dataset = build_ray_dataset_from_directory( + images_dir, + with_subdirs_as_labels=False, + remove_extension=False, + ) + # ray_dataset = ray_dataset.limit(100) # Limit to 100 images for testing + ray_dataset = filter_dataset_by_path(ray_dataset, paths_set) + predictions = run_object_detector_against_ray_dataset( + ray_dataset, + batch_size=64, + labels=labels, + inc_sub_labels_dict=inc_sub_labels_dict, + label_conf_thres=label_conf_thres, + nms_thre=0.9, + ) + _image_metadata, _categories_metadata, coco_format_predictions = ( + convert_detections_to_coco_format( + predictions, + labels, + images_metadata, + categories_metadata, + ) + ) + + # Save predictions in COCO format + with open(predictions_path, "w") as f: + json.dump(coco_format_predictions, f) + + # Load COCO predictions + coco_dt = coco_gt.loadRes(predictions_path) + + # Evaluate COCO predictions + coco_eval = COCOeval(coco_gt, coco_dt, "bbox") + coco_eval.params.imgIds = img_ids + coco_eval.evaluate() + coco_eval.accumulate() + coco_eval.summarize() diff --git a/directai_fastapi/classify_directory.py b/directai_fastapi/classify_directory.py index 5fb4498..9d5cf8a 100644 --- a/directai_fastapi/classify_directory.py +++ b/directai_fastapi/classify_directory.py @@ -3,7 +3,7 @@ from modeling.batch_processing import ( build_ray_dataset_from_directory, run_image_classifier_against_ray_dataset, - write_predictions_to_csv, + write_classifications_to_csv, filter_dataset_by_path, ) from pydantic_models import ClassifierDeploy @@ -52,7 +52,7 @@ def classify_directory( assert ( output_file is not None ), "Output file must be provided if not doing just an eval." - write_predictions_to_csv(predictions, output_file) + write_classifications_to_csv(predictions, output_file) if __name__ == "__main__": diff --git a/directai_fastapi/modeling/batch_processing.py b/directai_fastapi/modeling/batch_processing.py index de7fcb1..9656a6c 100644 --- a/directai_fastapi/modeling/batch_processing.py +++ b/directai_fastapi/modeling/batch_processing.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import ray from ray.data.datasource import PartitionStyle from ray.data.datasource.partitioning import Partitioning @@ -7,8 +9,15 @@ from torchvision.transforms import v2 # type: ignore[import-untyped] from torchvision.transforms.functional import InterpolationMode # type: ignore[import-untyped] import os +from typing import Any + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pycocotools.coco import _Image, _Category from modeling.image_classifier import ZeroShotImageClassifierWithFeedback +from modeling.object_detector import ZeroShotObjectDetectorWithFeedback def dir_name_to_label(dir_name: str) -> str: @@ -29,6 +38,7 @@ def truncate_full_path( def build_ray_dataset_from_directory( root: str, with_subdirs_as_labels: bool = True, + remove_extension: bool = True, ) -> ray.data.Dataset: # we expect images to be stored directly in the root directory # unless with_subdirs_as_labels is set to True @@ -49,7 +59,13 @@ def build_ray_dataset_from_directory( root, mode="RGB", partitioning=partitioning, include_paths=True ) - ds = ds.map(partial(truncate_full_path, last_n=2 if with_subdirs_as_labels else 1)) + ds = ds.map( + partial( + truncate_full_path, + last_n=2 if with_subdirs_as_labels else 1, + remove_extension=remove_extension, + ) + ) return ds @@ -106,6 +122,40 @@ def preprocess_image_for_classifier( return row +def preprocess_image_for_detector( + row: dict[str, np.ndarray], image_size: tuple[int, int] +) -> dict[str, np.ndarray]: + image = row["image"] + image_initial_size = image.shape[:2] + + padded_tensor = torch.ones((3, *image_size), dtype=torch.float32) + + r = min( + image_size[0] / image_initial_size[0], image_size[1] / image_initial_size[1] + ) + target_size = (int(image_initial_size[0] * r), int(image_initial_size[1] * r)) + + image = v2.functional.to_image(image) + image = v2.functional.to_dtype(image, torch.float32, scale=False) + image = v2.functional.resize( + image, target_size, interpolation=InterpolationMode.BICUBIC + ) + + assert isinstance(padded_tensor, torch.Tensor) + assert isinstance(image, torch.Tensor) + padded_tensor[:, : target_size[0], : target_size[1]] = image + + row["image"] = padded_tensor.numpy() + row["image_scale_ratio"] = np.array( + [ + r, + ] + ) + row["image_initial_size"] = np.array(image_initial_size) + + return row + + class RayDataImageClassifier: def __init__( self, @@ -124,19 +174,23 @@ def __init__( self.fill_cache() + def run_model(self, images: torch.Tensor) -> torch.Tensor: + with torch.inference_mode(), torch.autocast(str(self.model.device)): + raw_scores = self.model( + images, + labels=self.labels, + inc_sub_labels_dict=self.inc_sub_labels_dict, + exc_sub_labels_dict=self.exc_sub_labels_dict, + augment_examples=self.augment_examples, + ) + scores = torch.nn.functional.softmax(raw_scores / 0.07, dim=1) + return scores + def fill_cache(self) -> None: # compute embeddings for all prompts and warm up autocasting for _ in range(4): - rand_image = np.random.rand(1, 3, 224, 224).astype(np.float32) - batch: dict[str, np.ndarray] = { - "image": rand_image, - "label": np.array( - [ - "dummy", - ] - ), - } - self(batch) + rand_image = torch.rand((1, 3, 224, 224), device=self.device) + _ = self.run_model(rand_image) def __call__( self, @@ -144,17 +198,10 @@ def __call__( ) -> dict[str, np.ndarray]: images = torch.from_numpy(batch.pop("image")).to(self.device) - with torch.inference_mode(), torch.autocast(str(self.model.device)): - raw_scores = self.model( - images, - labels=self.labels, - inc_sub_labels_dict=self.inc_sub_labels_dict, - exc_sub_labels_dict=self.exc_sub_labels_dict, - augment_examples=self.augment_examples, - ) - scores = torch.nn.functional.softmax(raw_scores / 0.07, dim=1) - ind = scores.argmax(dim=1).cpu().numpy() - pred = np.array([self.labels[i] for i in ind]) + scores = self.run_model(images) + + ind = scores.argmax(dim=1).cpu().numpy() + pred = np.array([self.labels[i] for i in ind]) batch["scores"] = scores.cpu().numpy() batch["pred"] = pred @@ -169,6 +216,76 @@ def __call__( return batch +class RayDataObjectDetector: + def __init__( + self, + labels: list[str], + inc_sub_labels_dict: dict[str, list[str]], + exc_sub_labels_dict: dict[str, list[str]] | None = None, + label_conf_thres: dict[str, float] | None = None, + augment_examples: bool = True, + nms_thre: float = 0.4, + run_class_agnostic_nms: bool = True, + ) -> None: + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.model = ZeroShotObjectDetectorWithFeedback(device=self.device) + + self.labels = labels + self.inc_sub_labels_dict = inc_sub_labels_dict + self.exc_sub_labels_dict = exc_sub_labels_dict + self.label_conf_thres = label_conf_thres + self.augment_examples = augment_examples + self.nms_thre = nms_thre + self.run_class_agnostic_nms = run_class_agnostic_nms + + self.fill_cache() + + def run_model( + self, images: torch.Tensor, image_scale_ratios: torch.Tensor + ) -> list[list[torch.Tensor]]: + with torch.inference_mode(), torch.autocast(str(self.model.device)): + batched_predicted_boxes = self.model( + images, + labels=self.labels, + inc_sub_labels_dict=self.inc_sub_labels_dict, + exc_sub_labels_dict=self.exc_sub_labels_dict, + label_conf_thres=self.label_conf_thres, + augment_examples=self.augment_examples, + nms_thre=self.nms_thre, + run_class_agnostic_nms=self.run_class_agnostic_nms, + image_scale_ratios=image_scale_ratios, + ) + return batched_predicted_boxes + + def fill_cache(self) -> None: + # compute embeddings for all prompts and warm up autocasting + for _ in range(4): + rand_image = torch.rand((1, 3, 1008, 1008), device=self.device) + image_scale_ratio = torch.tensor( + [ + 1, + ], + device=self.device, + ) + _ = self.run_model(rand_image, image_scale_ratio) + + def __call__( + self, + batch: dict[str, np.ndarray], + ) -> dict[str, np.ndarray]: + images = torch.from_numpy(batch.pop("image")).to(self.device) + image_scale_ratios = torch.from_numpy(batch.pop("image_scale_ratio")).to( + self.device + ) + + batched_predicted_boxes = self.run_model(images, image_scale_ratios) + + # TODO: this may not be representable via a numpy array + batch["predicted_boxes"] = batched_predicted_boxes # type: ignore[assignment] + + return batch + + def run_image_classifier_against_ray_dataset( ds: ray.data.Dataset, batch_size: int, @@ -209,15 +326,61 @@ def run_image_classifier_against_ray_dataset( return predictions -def write_predictions_to_csv( - predictions: ray.data.Dataset, +def run_object_detector_against_ray_dataset( + ds: ray.data.Dataset, + batch_size: int, + labels: list[str], + inc_sub_labels_dict: dict[str, list[str]], + exc_sub_labels_dict: dict[str, list[str]] | None = None, + label_conf_thres: dict[str, float] | None = None, + augment_examples: bool = True, + nms_thre: float = 0.4, + run_class_agnostic_nms: bool = True, + concurrency: int | None = None, + num_gpus: int | None = None, +) -> ray.data.Dataset: + if num_gpus is None: + # the number of GPUs to reserve for each actor + num_gpus = 1 + + if concurrency is None: + # the number of actors to create + # set to the number of available GPUs + concurrency = torch.cuda.device_count() + + # we are going to assume that the object detector wants 1008x1008 images + preprocessed_ds = ds.map( + partial(preprocess_image_for_detector, image_size=(1008, 1008)) + ) + + predictions = preprocessed_ds.map_batches( + RayDataObjectDetector, + batch_size=batch_size, + concurrency=concurrency, + num_gpus=num_gpus, + fn_constructor_kwargs={ + "labels": labels, + "inc_sub_labels_dict": inc_sub_labels_dict, + "exc_sub_labels_dict": exc_sub_labels_dict, + "label_conf_thres": label_conf_thres, + "augment_examples": augment_examples, + "nms_thre": nms_thre, + "run_class_agnostic_nms": run_class_agnostic_nms, + }, + ) + + return predictions + + +def write_classifications_to_csv( + classifications: ray.data.Dataset, output_file: str, ) -> None: # Ray data has a built-in method to write to CSV # but it does it by partition into multiple files # we're just going to write everything to a single file - iterator = iter(predictions.iter_rows()) + iterator = iter(classifications.iter_rows()) with open(output_file, "w") as file: first_row = next(iterator) @@ -241,6 +404,106 @@ def write_row(row: dict[str, np.ndarray]) -> None: write_row(row) +def convert_detections_to_coco_format( + detections: ray.data.Dataset, + labels: list[str], + image_metadata: list[dict[str, Any]] | list["_Image"] | None = None, + categories_metadata: list[dict[str, Any]] | list["_Category"] | None = None, +) -> tuple[ + list[dict[str, Any]] | list["_Image"], + list[dict[str, Any]] | list["_Category"], + list[dict[str, Any]], +]: + # see https://docs.aws.amazon.com/rekognition/latest/customlabels-dg/md-coco-overview.html for the COCO format + # if image_metadata or categories_metadata are not provided, we will try to infer them from the detections + # NOTE: labels may not be redundant with categories_metadata, as the "id" field in categories_metadata may not be the same as the index of the label in labels + + need_to_infer_image_metadata = image_metadata is None + need_to_infer_categories_metadata = categories_metadata is None + + if need_to_infer_image_metadata: + image_metadata = [] + if need_to_infer_categories_metadata: + categories_metadata = [ + { + "id": i, + "name": label, + "supercategory": label, + } + for i, label in enumerate(labels) + ] + + assert isinstance(categories_metadata, list) + label_name_to_id = { + category["name"]: category["id"] for category in categories_metadata + } + label_name_to_ind = {label: i for i, label in enumerate(labels)} + assert label_name_to_ind.keys() == label_name_to_id.keys() + label_ind_to_id = { + label_name_to_ind[label]: label_name_to_id[label] for label in labels + } + + assert isinstance(image_metadata, list) + image_path_to_id = {data["file_name"]: data["id"] for data in image_metadata} + + annotations = [] + + for row in detections.iter_rows(): + predicted_boxes = row["predicted_boxes"] + image_h, image_w = [int(v) for v in row["image_initial_size"]] + + if row["path"] not in image_path_to_id: + assert ( + need_to_infer_image_metadata + ), f"Image metadata not provided for {row['path']}" + image_id = len(image_metadata) + image_metadata.append( + { + "id": image_id, + "file_name": row["path"], + "height": image_h, + "width": image_w, + } + ) + else: + image_id = image_path_to_id[row["path"]] + + for i, predicted_boxes_in_class in enumerate(predicted_boxes): + category_id = label_ind_to_id[i] + + for box in predicted_boxes_in_class: + x_1, y_1, x_2, y_2, score = box.tolist() + + x_1 = max(0, min(image_w, int(x_1))) + x_2 = max(0, min(image_w, int(x_2))) + y_1 = max(0, min(image_h, int(y_1))) + y_2 = max(0, min(image_h, int(y_2))) + + annotations.append( + { + "image_id": image_id, + "category_id": category_id, + "bbox": [x_1, y_1, x_2 - x_1, y_2 - y_1], + "score": float(score), + } + ) + + return image_metadata, categories_metadata, annotations + + +def build_naive_detector_config_from_coco_categories_metadata( + categories_metadata: list[dict[str, Any]] | list["_Category"], +) -> tuple[list[str], dict[str, list[str]]]: + labels = [category["name"] for category in categories_metadata] + inc_sub_labels_dict = { + label: [ + label, + ] + for label in labels + } + return labels, inc_sub_labels_dict + + if __name__ == "__main__": # first we load the pytorch dataset and save it to disk # we're not going to actually use it, but we want to make sure it's available @@ -279,7 +542,7 @@ def write_row(row: dict[str, np.ndarray]) -> None: augment_examples=augment_examples, ) - write_predictions_to_csv( + write_classifications_to_csv( predictions, os.path.join(food101_root_dir, "food-101", "meta", f"{split}_predictions.csv"), ) diff --git a/directai_fastapi/modeling/object_detector.py b/directai_fastapi/modeling/object_detector.py index b29274a..a6af3b3 100644 --- a/directai_fastapi/modeling/object_detector.py +++ b/directai_fastapi/modeling/object_detector.py @@ -551,6 +551,8 @@ def run_nms_via_adjacency_list( valid_box_indices_by_descending_score: torch.Tensor, adjacency_list: list[torch.Tensor], ) -> torch.Tensor: + if valid_box_indices_by_descending_score.numel() == 0: + return valid_box_indices_by_descending_score # we compute the indices of the start and end of each box's adjacent boxes # since our graph representation is just a list of edges, we would like to know which edges correspond to which nodes # as the first node is sorted, we can just take the difference between adjacent nodes to get the start and end of each node's edges diff --git a/directai_fastapi/requirements.txt b/directai_fastapi/requirements.txt index 3cc6bd6..e9f8a0b 100644 --- a/directai_fastapi/requirements.txt +++ b/directai_fastapi/requirements.txt @@ -14,4 +14,8 @@ https://data.pyg.org/whl/torch-2.2.0%2Bcu121/torch_scatter-2.1.2%2Bpt22cu121-cp3 lru-dict==1.3.0 transformers==4.44.2 flash-attn==2.6.3 -fire==0.6.0 \ No newline at end of file +fire==0.6.0 +cython==3.0.11 +cython-bbox==0.1.5 +pycocotools==2.0.8 +types-pycocotools==2.0.0.20240806 \ No newline at end of file diff --git a/directai_fastapi/run.sh b/directai_fastapi/run.sh index 734b011..6133332 100644 --- a/directai_fastapi/run.sh +++ b/directai_fastapi/run.sh @@ -1,5 +1,6 @@ timestamp=$(date +%s) log_dir="logs/fastapi_$timestamp" mkdir -p "$log_dir" -python -m unittest unit_tests/test.py -uvicorn server:app --host 0.0.0.0 --port 8000 --log-level warning \ No newline at end of file +# python -m unittest unit_tests/test.py +# uvicorn server:app --host 0.0.0.0 --port 8000 --log-level warning +python benchmarking/benchmark_coco.py \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml index 5e0d50d..174e771 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,14 +1,14 @@ version: '2.3' services: - local_redis: - build: redis_data/ - container_name: local_redis - ports: - - 6379:6379 - volumes: - - ./redis_data/:/data - networks: - - deploy_network + # local_redis: + # build: redis_data/ + # container_name: local_redis + # ports: + # - 6379:6379 + # volumes: + # - ./redis_data/:/data + # networks: + # - deploy_network local_fastapi: build: directai_fastapi/ ports: @@ -18,7 +18,7 @@ services: container_name: local_fastapi environment: - PYTHONUNBUFFERED=1 - - NVIDIA_VISIBLE_DEVICES=all + - NVIDIA_VISIBLE_DEVICES=1 - HF_HOME=/directai_fastapi/.cache/huggingface - CACHE_REDIS_PORT=6379 env_file: @@ -28,23 +28,23 @@ services: - ./logs:/directai_fastapi/logs - ./.cache:/directai_fastapi/.cache shm_size: 10.24g # because Ray complains if it's less - depends_on: - - local_redis - extra_hosts: - - "host.docker.internal:host-gateway" - local_gradio: - build: gradio/ - networks: - - deploy_network - ports: - - 7860:7860 - depends_on: - - local_fastapi - container_name: local_gradio - environment: - - PYTHONUNBUFFERED=1 + # depends_on: + # - local_redis extra_hosts: - "host.docker.internal:host-gateway" + # local_gradio: + # build: gradio/ + # networks: + # - deploy_network + # ports: + # - 7860:7860 + # depends_on: + # - local_fastapi + # container_name: local_gradio + # environment: + # - PYTHONUNBUFFERED=1 + # extra_hosts: + # - "host.docker.internal:host-gateway" networks: deploy_network: