From 56bce9c1c8492910f350aeb9d69a3b44755abc3a Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Thu, 23 Sep 2021 10:56:01 -0400
Subject: [PATCH 01/35] .

---
 CHANGELOG.md                           |  2 ++
 flash/core/data/utils.py               | 14 ++++++++++++++
 flash/core/utilities/imports.py        |  1 +
 requirements/datatype_image_extras.txt |  1 +
 4 files changed, 18 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 11d4b0accf..3b95dc2ba7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added support for `from_data_frame` to `TextClassificationData` ([#785](https://github.com/PyTorchLightning/lightning-flash/pull/785))
 
+- Added `FastFace` integration ([#606](https://github.com/PyTorchLightning/lightning-flash/pull/606))
+
 ### Changed
 
 - Changed the default `num_workers` on linux to `0` (matching the default for other OS) ([#759](https://github.com/PyTorchLightning/lightning-flash/pull/759))
diff --git a/flash/core/data/utils.py b/flash/core/data/utils.py
index 3779b7426e..01b0f2c007 100644
--- a/flash/core/data/utils.py
+++ b/flash/core/data/utils.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os.path
+import tarfile
 import zipfile
 from typing import Any, Callable, Dict, Iterable, Mapping, Optional, Set, Type
 
@@ -148,10 +149,23 @@ def download_data(url: str, path: str = "data/", verbose: bool = False) -> None:
             ):
                 fp.write(chunk)  # type: ignore
 
+    def extract_tarfile(file_path: str, extract_path: str, mode: str):
+        if os.path.exists(file_path):
+            with tarfile.open(file_path, mode=mode) as tar_ref:
+                for member in tar_ref.getmembers():
+                    try:
+                        tar_ref.extract(member, path=extract_path, set_attrs=False)
+                    except PermissionError:
+                        pass
+
     if ".zip" in local_filename:
         if os.path.exists(local_filename):
             with zipfile.ZipFile(local_filename, "r") as zip_ref:
                 zip_ref.extractall(path)
+    elif local_filename.endswith(".tar.gz") or local_filename.endswith(".tgz"):
+        extract_tarfile(local_filename, path, "r:gz")
+    elif local_filename.endswith('.tar.bz2') or local_filename.endswith('.tbz'):
+        extract_tarfile(local_filename, path, "r:bz2")
 
 
 def _contains_any_tensor(value: Any, dtype: Type = Tensor) -> bool:
diff --git a/flash/core/utilities/imports.py b/flash/core/utilities/imports.py
index f138eaf37e..95a6272072 100644
--- a/flash/core/utilities/imports.py
+++ b/flash/core/utilities/imports.py
@@ -87,6 +87,7 @@ def _compare_version(package: str, op, version) -> bool:
 _PIL_AVAILABLE = _module_available("PIL")
 _OPEN3D_AVAILABLE = _module_available("open3d")
 _SEGMENTATION_MODELS_AVAILABLE = _module_available("segmentation_models_pytorch")
+_FASTFACE_AVAILABLE = _module_available("fastface")
 _LIBROSA_AVAILABLE = _module_available("librosa")
 _TORCH_SCATTER_AVAILABLE = _module_available("torch_scatter")
 _TORCH_SPARSE_AVAILABLE = _module_available("torch_sparse")
diff --git a/requirements/datatype_image_extras.txt b/requirements/datatype_image_extras.txt
index 7704ebb7a2..8e976d2ade 100644
--- a/requirements/datatype_image_extras.txt
+++ b/requirements/datatype_image_extras.txt
@@ -8,3 +8,4 @@ effdet
 albumentations
 learn2learn
 baal
+fastface

From 2b22d87d8fc366513e6f617c95c676fdeae45b1c Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Thu, 23 Sep 2021 12:16:37 -0400
Subject: [PATCH 02/35] merging taskathon PR code

---
 flash/image/__init__.py                |   1 +
 flash/image/face_detection/__init__.py |   1 +
 flash/image/face_detection/data.py     |  86 +++++++++
 flash/image/face_detection/model.py    | 236 +++++++++++++++++++++++++
 flash_examples/face_detection.py       |  50 ++++++
 5 files changed, 374 insertions(+)
 create mode 100644 flash/image/face_detection/__init__.py
 create mode 100644 flash/image/face_detection/data.py
 create mode 100644 flash/image/face_detection/model.py
 create mode 100644 flash_examples/face_detection.py

diff --git a/flash/image/__init__.py b/flash/image/__init__.py
index b3ac7f10b6..d881de0e70 100644
--- a/flash/image/__init__.py
+++ b/flash/image/__init__.py
@@ -6,6 +6,7 @@
 from flash.image.classification.backbones import IMAGE_CLASSIFIER_BACKBONES  # noqa: F401
 from flash.image.detection import ObjectDetectionData, ObjectDetector  # noqa: F401
 from flash.image.embedding import ImageEmbedder  # noqa: F401
+from flash.image.face_detection import FaceDetector # noqa: F401
 from flash.image.instance_segmentation import InstanceSegmentation, InstanceSegmentationData  # noqa: F401
 from flash.image.keypoint_detection import KeypointDetectionData, KeypointDetector  # noqa: F401
 from flash.image.segmentation import (  # noqa: F401
diff --git a/flash/image/face_detection/__init__.py b/flash/image/face_detection/__init__.py
new file mode 100644
index 0000000000..c642f1c2ba
--- /dev/null
+++ b/flash/image/face_detection/__init__.py
@@ -0,0 +1 @@
+from flash.image.face_detection.model import FaceDetector # noqa: F401
diff --git a/flash/image/face_detection/data.py b/flash/image/face_detection/data.py
new file mode 100644
index 0000000000..be8e8e9cd8
--- /dev/null
+++ b/flash/image/face_detection/data.py
@@ -0,0 +1,86 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, Callable, Dict, Optional, Sequence, Tuple
+
+from torch.utils.data import Dataset
+
+from flash.core.data.data_source import DataSource, DefaultDataKeys, DefaultDataSources
+from flash.core.data.process import Preprocess
+from flash.core.utilities.imports import _TORCHVISION_AVAILABLE
+from flash.image.data import ImagePathsDataSource
+from flash.image.detection.transforms import default_transforms
+
+if _TORCHVISION_AVAILABLE:
+    from torchvision.datasets.folder import default_loader
+
+
+class FastFaceDataSource(DataSource[Tuple[str, str]]):
+
+    def load_data(self, data: Dataset, dataset: Optional[Any] = None) -> Sequence[Dict[str, Any]]:
+        new_data = []
+        for img_file_path, targets in zip(data.ids, data.targets):
+            new_data.append(
+                dict(
+                    input=img_file_path,
+                    target=dict(
+                        boxes=targets["target_boxes"],
+                        labels=[1 for _ in range(targets["target_boxes"].shape[0])],
+                    )
+                )
+            )
+        return new_data
+
+    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
+        filepath = sample[DefaultDataKeys.INPUT]
+        img = default_loader(filepath)
+        sample[DefaultDataKeys.INPUT] = img
+        w, h = img.size  # WxH
+        sample[DefaultDataKeys.METADATA] = {
+            "filepath": filepath,
+            "size": (h, w),
+        }
+        return sample
+
+
+class FaceDetectionPreprocess(Preprocess):
+
+    def __init__(
+        self,
+        train_transform: Optional[Dict[str, Callable]] = None,
+        val_transform: Optional[Dict[str, Callable]] = None,
+        test_transform: Optional[Dict[str, Callable]] = None,
+        predict_transform: Optional[Dict[str, Callable]] = None
+    ):
+        super().__init__(
+            train_transform=train_transform,
+            val_transform=val_transform,
+            test_transform=test_transform,
+            predict_transform=predict_transform,
+            data_sources={
+                DefaultDataSources.FILES: ImagePathsDataSource(),
+                DefaultDataSources.FOLDERS: ImagePathsDataSource(),
+                "fastface": FastFaceDataSource()
+            },
+            default_data_source=DefaultDataSources.FILES,
+        )
+
+    def get_state_dict(self) -> Dict[str, Any]:
+        return {**self.transforms}
+
+    @classmethod
+    def load_state_dict(cls, state_dict: Dict[str, Any], strict: bool = False):
+        return cls(**state_dict)
+
+    def default_transforms(self) -> Optional[Dict[str, Callable]]:
+        return default_transforms()
diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
new file mode 100644
index 0000000000..a36d52de3c
--- /dev/null
+++ b/flash/image/face_detection/model.py
@@ -0,0 +1,236 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, Callable, List, Mapping, Optional, Sequence, Type, Union
+
+import torch
+from torch import nn
+from torch.optim import Optimizer
+
+from flash.core.data.data_source import DefaultDataKeys
+from flash.core.data.process import Preprocess, Serializer
+from flash.core.model import Task
+from flash.core.utilities.imports import _FASTFACE_AVAILABLE
+from flash.image.detection.finetuning import ObjectDetectionFineTuning
+from flash.image.detection.serialization import DetectionLabels
+from flash.image.face_detection.data import FaceDetectionPreprocess
+
+if _FASTFACE_AVAILABLE:
+    import fastface as ff
+
+
+class FaceDetector(Task):
+    """The ``FaceDetector`` is a :class:`~flash.Task` for detecting faces in images. For more details, see
+    :ref:`face_detection`.
+    Args:
+        model: a string of :attr`_models`. Defaults to 'lffd_slim'.
+        pretrained: Whether the model from fastface should be loaded with it's pretrained weights.
+        loss: the function(s) to update the model with. Has no effect for fastface models.
+        metrics: The provided metrics. All metrics here will be logged to progress bar and the respective logger.
+            Changing this argument currently has no effect.
+        optimizer: The optimizer to use for training. Can either be the actual class or the class name.
+        learning_rate: The learning rate to use for training
+    """
+
+    required_extras: str = "image"
+
+    def __init__(
+        self,
+        model: str = "lffd_slim",
+        pretrained: bool = True,
+        loss=None,
+        metrics: Union[Callable, nn.Module, Mapping, Sequence, None] = None,
+        optimizer: Type[Optimizer] = torch.optim.AdamW,
+        learning_rate: float = 1e-4,
+        serializer: Optional[Union[Serializer, Mapping[str, Serializer]]] = None,
+        preprocess: Optional[Preprocess] = None,
+        **kwargs: Any,
+    ):
+        self.save_hyperparameters()
+
+        if model in ff.list_pretrained_models():
+            model = FaceDetector.get_model(model, pretrained, **kwargs)
+        else:
+            ValueError(f"{model} is not supported yet.")
+
+        super().__init__(
+            model=model,
+            loss_fn=loss,
+            metrics=metrics or {"AP": ff.metric.AveragePrecision()},
+            learning_rate=learning_rate,
+            optimizer=optimizer,
+            serializer=serializer or DetectionLabels(),
+            preprocess=preprocess or FaceDetectionPreprocess(),
+        )
+
+    @staticmethod
+    def get_model(
+        model_name,
+        pretrained,
+        **kwargs,
+    ):
+
+        if pretrained:
+            pl_model = ff.FaceDetector.from_pretrained(model_name, **kwargs)
+        else:
+            arch, config = model_name.split("_")
+            pl_model = ff.FaceDetector.build(arch, config, **kwargs)
+
+        # get torch.nn.Module
+        model = getattr(pl_model, "arch")
+
+        # set preprocess params
+        model.register_buffer("normalizer", getattr(pl_model, "normalizer"))
+        model.register_buffer("mean", getattr(pl_model, "mean"))
+        model.register_buffer("std", getattr(pl_model, "std"))
+
+        # set postprocess function
+        setattr(model, "_postprocess", getattr(pl_model, "_postprocess"))
+
+        return model
+
+    def forward(self, x: List[torch.Tensor]) -> Any:
+
+        batch, scales, paddings = ff.utils.preprocess.prepare_batch(x, None, adaptive_batch=True)
+        # batch: torch.Tensor(B,C,T,T)
+        # scales: torch.Tensor(B,)
+        # paddings: torch.Tensor(B,4) as pad (left, top, right, bottom)
+
+        # apply preprocess
+        batch = (((batch * 255) / self.model.normalizer) - self.model.mean) / self.model.std
+
+        # get logits
+        logits = self.model(batch)
+        # logits, any
+
+        preds = self.model.logits_to_preds(logits)
+        # preds: torch.Tensor(B, N, 5)
+
+        preds = self.model._postprocess(preds)
+        # preds: torch.Tensor(N, 6) as x1,y1,x2,y2,score,batch_idx
+
+        preds = [preds[preds[:, 5] == batch_idx, :5] for batch_idx in range(batch.size(0))]
+        # preds: list of torch.Tensor(N, 5) as x1,y1,x2,y2,score
+
+        preds = ff.utils.preprocess.adjust_results(preds, scales, paddings)
+        # preds: list of torch.Tensor(N, 5) as x1,y1,x2,y2,score
+
+        return preds
+
+    def _prepare_batch(self, batch):
+        images, targets = batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET]
+
+        targets = [{"target_boxes": target["boxes"]} for target in targets]
+
+        batch, scales, paddings = ff.utils.preprocess.prepare_batch(images, None, adaptive_batch=True)
+        # batch: torch.Tensor(B,C,T,T)
+        # scales: torch.Tensor(B,)
+        # paddings: torch.Tensor(B,4) as pad (left, top, right, bottom)
+
+        # apply preprocess
+        batch = (((batch * 255) / self.model.normalizer) - self.model.mean) / self.model.std
+
+        # adjust targets
+        for i, (target, scale, padding) in enumerate(zip(targets, scales, paddings)):
+            target["target_boxes"] *= scale
+            target["target_boxes"][:, [0, 2]] += padding[0]
+            target["target_boxes"][:, [1, 3]] += padding[1]
+            targets[i]["target_boxes"] = target["target_boxes"]
+
+        return batch, targets
+
+    def _compute_metrics(self, logits, targets):
+        preds = self.model.logits_to_preds(logits)
+        # preds: torch.Tensor(B, N, 5)
+
+        preds = self.model._postprocess(preds)
+        # preds: torch.Tensor(N, 6) as x1,y1,x2,y2,score,batch_idx
+
+        target_boxes = [target["target_boxes"] for target in targets]
+        pred_boxes = [preds[preds[:, 5] == batch_idx, :5] for batch_idx in range(len(targets))]
+
+        for metric in self.val_metrics.values():
+            metric.update(pred_boxes, target_boxes)
+
+    def training_step(self, batch, batch_idx) -> Any:
+        """The training step. Overrides ``Task.training_step``
+        """
+
+        batch, targets = self._prepare_batch(batch)
+
+        # get logits
+        logits = self.model(batch)
+        # logits, any
+
+        # compute loss
+        loss = self.model.compute_loss(logits, targets)
+        # loss: dict of losses or loss
+
+        self.log_dict({f"train_{k}": v for k, v in loss.items()}, on_step=True, on_epoch=True, prog_bar=True)
+        return loss
+
+    def on_validation_epoch_start(self) -> None:
+        for metric in self.val_metrics.values():
+            metric.reset()
+
+    def validation_step(self, batch, batch_idx):
+        batch, targets = self._prepare_batch(batch)
+
+        # get logits
+        logits = self.model(batch)
+        # logits, any
+
+        # compute loss
+        loss = self.model.compute_loss(logits, targets)
+        # loss: dict of losses or loss
+
+        self._compute_metrics(logits, targets)
+
+        self.log_dict({f"val_{k}": v for k, v in loss.items()}, on_step=True, on_epoch=True, prog_bar=True)
+        return loss
+
+    def validation_epoch_end(self, outputs) -> None:
+        metric_results = {name: metric.compute() for name, metric in self.val_metrics.items()}
+        self.log_dict({f"val_{k}": v for k, v in metric_results.items()}, on_epoch=True)
+
+    def on_test_epoch_start(self) -> None:
+        for metric in self.val_metrics.values():
+            metric.reset()
+
+    def test_step(self, batch, batch_idx):
+        batch, targets = self._prepare_batch(batch)
+
+        # get logits
+        logits = self.model(batch)
+        # logits, any
+
+        # compute loss
+        loss = self.model.compute_loss(logits, targets)
+        # loss: dict of losses or loss
+
+        self._compute_metrics(logits, targets)
+
+        self.log_dict({f"test_{k}": v for k, v in loss.items()}, on_step=True, on_epoch=True, prog_bar=True)
+        return loss
+
+    def test_epoch_end(self, outputs) -> None:
+        metric_results = {name: metric.compute() for name, metric in self.val_metrics.items()}
+        self.log_dict({f"test_{k}": v for k, v in metric_results.items()}, on_epoch=True)
+
+    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
+        images = batch[DefaultDataKeys.INPUT]
+        batch[DefaultDataKeys.PREDS] = self(images)
+        return batch
+
+    def configure_finetune_callback(self):
+        return [ObjectDetectionFineTuning(train_bn=True)]
diff --git a/flash_examples/face_detection.py b/flash_examples/face_detection.py
new file mode 100644
index 0000000000..e8a14f7096
--- /dev/null
+++ b/flash_examples/face_detection.py
@@ -0,0 +1,50 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import flash
+from flash.core.data.data_module import DataModule
+from flash.core.utilities.imports import _FASTFACE_AVAILABLE
+from flash.image import FaceDetector
+from flash.image.face_detection.data import FaceDetectionPreprocess
+
+if _FASTFACE_AVAILABLE:
+    import fastface as ff
+else:
+    raise ModuleNotFoundError("Please, pip install -e '.[image]'")
+
+# 1. Create the DataModule
+train_dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="train")
+val_dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="val")
+
+datamodule = DataModule.from_data_source(
+    "fastface", train_data=train_dataset, val_data=val_dataset, preprocess=FaceDetectionPreprocess()
+)
+
+# 2. Build the task
+model = FaceDetector(model="lffd_slim")
+
+# 3. Create the trainer and finetune the model
+trainer = flash.Trainer(max_epochs=3, limit_train_batches=0.1, limit_val_batches=0.1)
+
+trainer.finetune(model, datamodule=datamodule, strategy="freeze")
+
+# 4. Detect faces in a few images!
+predictions = model.predict([
+    "data/2002/07/19/big/img_18.jpg",
+    "data/2002/07/19/big/img_65.jpg",
+    "data/2002/07/19/big/img_255.jpg",
+])
+print(predictions)
+
+# 5. Save the model!
+trainer.save_checkpoint("face_detection_model.pt")

From 6c624ad847dc07c73eedd2143522eec263634a6f Mon Sep 17 00:00:00 2001
From: Ananya Harsh Jha <ananya@pytorchlightning.ai>
Date: Fri, 24 Sep 2021 07:34:34 -0400
Subject: [PATCH 03/35] working

---
 flash/image/__init__.py                |   2 +-
 flash/image/face_detection/__init__.py |   3 +-
 flash/image/face_detection/data.py     | 118 ++++++++++++++++++++----
 flash/image/face_detection/model.py    | 119 +++++++++----------------
 flash_examples/face_detection.py       |  20 +++--
 5 files changed, 156 insertions(+), 106 deletions(-)

diff --git a/flash/image/__init__.py b/flash/image/__init__.py
index d881de0e70..36257c77ce 100644
--- a/flash/image/__init__.py
+++ b/flash/image/__init__.py
@@ -6,7 +6,7 @@
 from flash.image.classification.backbones import IMAGE_CLASSIFIER_BACKBONES  # noqa: F401
 from flash.image.detection import ObjectDetectionData, ObjectDetector  # noqa: F401
 from flash.image.embedding import ImageEmbedder  # noqa: F401
-from flash.image.face_detection import FaceDetector # noqa: F401
+from flash.image.face_detection import FaceDetector, FaceDetectionData # noqa: F401
 from flash.image.instance_segmentation import InstanceSegmentation, InstanceSegmentationData  # noqa: F401
 from flash.image.keypoint_detection import KeypointDetectionData, KeypointDetector  # noqa: F401
 from flash.image.segmentation import (  # noqa: F401
diff --git a/flash/image/face_detection/__init__.py b/flash/image/face_detection/__init__.py
index c642f1c2ba..7d14bba121 100644
--- a/flash/image/face_detection/__init__.py
+++ b/flash/image/face_detection/__init__.py
@@ -1 +1,2 @@
-from flash.image.face_detection.model import FaceDetector # noqa: F401
+from flash.image.face_detection.model import FaceDetector  # noqa: F401
+from flash.image.face_detection.data import FaceDetectionData  # noqa: F401
diff --git a/flash/image/face_detection/data.py b/flash/image/face_detection/data.py
index be8e8e9cd8..8d196819e8 100644
--- a/flash/image/face_detection/data.py
+++ b/flash/image/face_detection/data.py
@@ -11,45 +11,90 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, Dict, Optional, Sequence, Tuple
+from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Mapping
+
+import torch
+import torchvision
+import torch.nn as nn
 
 from torch.utils.data import Dataset
 
-from flash.core.data.data_source import DataSource, DefaultDataKeys, DefaultDataSources
-from flash.core.data.process import Preprocess
-from flash.core.utilities.imports import _TORCHVISION_AVAILABLE
+from flash.core.data.transforms import ApplyToKeys
+from flash.core.data.data_source import DatasetDataSource, DefaultDataKeys, DefaultDataSources
+from flash.core.data.process import Preprocess, Postprocess
+from flash.core.utilities.imports import _TORCHVISION_AVAILABLE, _ICEVISION_AVAILABLE, _FASTFACE_AVAILABLE
 from flash.image.data import ImagePathsDataSource
-from flash.image.detection.transforms import default_transforms
+from flash.core.integrations.icevision.data import IceVisionParserDataSource
+from flash.core.integrations.icevision.transforms import default_transforms
+from flash.image.detection import ObjectDetectionData
 
 if _TORCHVISION_AVAILABLE:
     from torchvision.datasets.folder import default_loader
 
+if _ICEVISION_AVAILABLE:
+    from icevision.parsers import COCOBBoxParser
+else:
+    COCOBBoxParser = object
+
+if _FASTFACE_AVAILABLE:
+    import fastface as ff
+
+
+def fastface_collate_fn(samples: Sequence[Dict[str, Any]]) -> Dict[str, Sequence[Any]]:
+    samples = {key: [sample[key] for sample in samples] for key in samples[0]}
+
+    images, scales, paddings = ff.utils.preprocess.prepare_batch(
+        samples[DefaultDataKeys.INPUT], None, adaptive_batch=True
+    )
+
+    samples["scales"] = scales
+    samples["paddings"] = paddings
 
-class FastFaceDataSource(DataSource[Tuple[str, str]]):
+    if DefaultDataKeys.TARGET in samples.keys():
+        targets = samples[DefaultDataKeys.TARGET]
+        targets = [{"target_boxes": target["boxes"]} for target in targets]
 
-    def load_data(self, data: Dataset, dataset: Optional[Any] = None) -> Sequence[Dict[str, Any]]:
+        for i, (target, scale, padding) in enumerate(zip(targets, scales, paddings)):
+            target["target_boxes"] *= scale
+            target["target_boxes"][:, [0, 2]] += padding[0]
+            target["target_boxes"][:, [1, 3]] += padding[1]
+            targets[i]["target_boxes"] = target["target_boxes"]
+
+        samples[DefaultDataKeys.TARGET] = targets
+    samples[DefaultDataKeys.INPUT] = images
+
+    return samples
+
+
+class FastFaceDataSource(DatasetDataSource):
+    def load_data(self, data: Dataset, dataset: Any = None) -> Dataset:
         new_data = []
         for img_file_path, targets in zip(data.ids, data.targets):
             new_data.append(
-                dict(
-                    input=img_file_path,
-                    target=dict(
-                        boxes=targets["target_boxes"],
-                        labels=[1 for _ in range(targets["target_boxes"].shape[0])],
+                super().load_sample(
+                    (
+                        img_file_path,
+                        dict(
+                            boxes=targets["target_boxes"],
+                            labels=[1 for _ in range(targets["target_boxes"].shape[0])],
+                        )
                     )
                 )
             )
+
         return new_data
 
-    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
+    def load_sample(self, sample: Any, dataset: Optional[Any] = None) -> Mapping[str, Any]:
         filepath = sample[DefaultDataKeys.INPUT]
         img = default_loader(filepath)
         sample[DefaultDataKeys.INPUT] = img
+
         w, h = img.size  # WxH
         sample[DefaultDataKeys.METADATA] = {
             "filepath": filepath,
             "size": (h, w),
         }
+
         return sample
 
 
@@ -60,8 +105,11 @@ def __init__(
         train_transform: Optional[Dict[str, Callable]] = None,
         val_transform: Optional[Dict[str, Callable]] = None,
         test_transform: Optional[Dict[str, Callable]] = None,
-        predict_transform: Optional[Dict[str, Callable]] = None
+        predict_transform: Optional[Dict[str, Callable]] = None,
+        image_size: Tuple[int, int] = (128, 128),
     ):
+        self.image_size = image_size
+
         super().__init__(
             train_transform=train_transform,
             val_transform=val_transform,
@@ -70,7 +118,7 @@ def __init__(
             data_sources={
                 DefaultDataSources.FILES: ImagePathsDataSource(),
                 DefaultDataSources.FOLDERS: ImagePathsDataSource(),
-                "fastface": FastFaceDataSource()
+                DefaultDataSources.DATASETS: FastFaceDataSource(),
             },
             default_data_source=DefaultDataSources.FILES,
         )
@@ -82,5 +130,41 @@ def get_state_dict(self) -> Dict[str, Any]:
     def load_state_dict(cls, state_dict: Dict[str, Any], strict: bool = False):
         return cls(**state_dict)
 
-    def default_transforms(self) -> Optional[Dict[str, Callable]]:
-        return default_transforms()
+    def default_transforms(self) -> Dict[str, Callable]:
+        return {
+            "to_tensor_transform": nn.Sequential(
+                ApplyToKeys(DefaultDataKeys.INPUT, torchvision.transforms.ToTensor()),
+                ApplyToKeys(
+                    DefaultDataKeys.TARGET,
+                    nn.Sequential(
+                        ApplyToKeys('boxes', torch.as_tensor),
+                        ApplyToKeys('labels', torch.as_tensor),
+                    )
+                ),
+            ),
+            "collate": fastface_collate_fn,
+        }
+
+
+class FaceDetectionPostProcess(Postprocess):
+    @staticmethod
+    def per_batch_transform(batch: Any) -> Any:
+        scales = batch['scales']
+        paddings = batch['paddings']
+
+        batch.pop('scales', None)
+        batch.pop('paddings', None)
+
+        preds = batch[DefaultDataKeys.PREDS]
+
+        # preds: list of torch.Tensor(N, 5) as x1, y1, x2, y2, score
+        preds = [preds[preds[:, 5] == batch_idx, :5] for batch_idx in range(len(preds))]
+        preds = ff.utils.preprocess.adjust_results(preds, scales, paddings)
+        batch[DefaultDataKeys.PREDS] = preds
+
+        return batch
+
+
+class FaceDetectionData(ObjectDetectionData):
+    preprocess_cls = FaceDetectionPreprocess
+    postprocess_cls = FaceDetectionPostProcess
diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index a36d52de3c..0679247935 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -11,24 +11,42 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, List, Mapping, Optional, Sequence, Type, Union
+from typing import Any, Callable, List, Mapping, Optional, Sequence, Type, Union, Dict
 
 import torch
+import pytorch_lightning as pl
+
 from torch import nn
 from torch.optim import Optimizer
 
+from flash.core.model import Task
 from flash.core.data.data_source import DefaultDataKeys
+from flash.core.data.process import Postprocess
 from flash.core.data.process import Preprocess, Serializer
-from flash.core.model import Task
 from flash.core.utilities.imports import _FASTFACE_AVAILABLE
-from flash.image.detection.finetuning import ObjectDetectionFineTuning
-from flash.image.detection.serialization import DetectionLabels
-from flash.image.face_detection.data import FaceDetectionPreprocess
+from flash.core.finetuning import FlashBaseFinetuning
+from flash.image.face_detection.data import FaceDetectionPreprocess, FaceDetectionPostProcess
 
 if _FASTFACE_AVAILABLE:
     import fastface as ff
 
 
+class FaceDetectionFineTuning(FlashBaseFinetuning):
+    def __init__(self, train_bn: bool = True) -> None:
+        super().__init__(train_bn=train_bn)
+
+    def freeze_before_training(self, pl_module: pl.LightningModule) -> None:
+        self.freeze(modules=pl_module.model.backbone, train_bn=self.train_bn)
+
+
+class DetectionLabels(Serializer):
+    """A :class:`.Serializer` which extracts predictions from sample dict."""
+
+    def serialize(self, sample: Any) -> Dict[str, Any]:
+        sample = sample[DefaultDataKeys.PREDS] if isinstance(sample, Dict) else sample
+        return sample
+
+
 class FaceDetector(Task):
     """The ``FaceDetector`` is a :class:`~flash.Task` for detecting faces in images. For more details, see
     :ref:`face_detection`.
@@ -100,54 +118,19 @@ def get_model(
         return model
 
     def forward(self, x: List[torch.Tensor]) -> Any:
+        images = self._prepare_batch(x)
+        logits = self.model(images)
 
-        batch, scales, paddings = ff.utils.preprocess.prepare_batch(x, None, adaptive_batch=True)
-        # batch: torch.Tensor(B,C,T,T)
-        # scales: torch.Tensor(B,)
-        # paddings: torch.Tensor(B,4) as pad (left, top, right, bottom)
-
-        # apply preprocess
-        batch = (((batch * 255) / self.model.normalizer) - self.model.mean) / self.model.std
-
-        # get logits
-        logits = self.model(batch)
-        # logits, any
-
-        preds = self.model.logits_to_preds(logits)
         # preds: torch.Tensor(B, N, 5)
-
-        preds = self.model._postprocess(preds)
         # preds: torch.Tensor(N, 6) as x1,y1,x2,y2,score,batch_idx
-
-        preds = [preds[preds[:, 5] == batch_idx, :5] for batch_idx in range(batch.size(0))]
-        # preds: list of torch.Tensor(N, 5) as x1,y1,x2,y2,score
-
-        preds = ff.utils.preprocess.adjust_results(preds, scales, paddings)
-        # preds: list of torch.Tensor(N, 5) as x1,y1,x2,y2,score
+        preds = self.model.logits_to_preds(logits)
+        preds = self.model._postprocess(preds)
 
         return preds
 
     def _prepare_batch(self, batch):
-        images, targets = batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET]
-
-        targets = [{"target_boxes": target["boxes"]} for target in targets]
-
-        batch, scales, paddings = ff.utils.preprocess.prepare_batch(images, None, adaptive_batch=True)
-        # batch: torch.Tensor(B,C,T,T)
-        # scales: torch.Tensor(B,)
-        # paddings: torch.Tensor(B,4) as pad (left, top, right, bottom)
-
-        # apply preprocess
         batch = (((batch * 255) / self.model.normalizer) - self.model.mean) / self.model.std
-
-        # adjust targets
-        for i, (target, scale, padding) in enumerate(zip(targets, scales, paddings)):
-            target["target_boxes"] *= scale
-            target["target_boxes"][:, [0, 2]] += padding[0]
-            target["target_boxes"][:, [1, 3]] += padding[1]
-            targets[i]["target_boxes"] = target["target_boxes"]
-
-        return batch, targets
+        return batch
 
     def _compute_metrics(self, logits, targets):
         preds = self.model.logits_to_preds(logits)
@@ -162,19 +145,19 @@ def _compute_metrics(self, logits, targets):
         for metric in self.val_metrics.values():
             metric.update(pred_boxes, target_boxes)
 
-    def training_step(self, batch, batch_idx) -> Any:
-        """The training step. Overrides ``Task.training_step``
-        """
+    def shared_step(self, batch, train=False) -> Any:
+        images, targets = batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET]
+        images = self._prepare_batch(images)
+        logits = self.model(images)
+        loss = self.model.compute_loss(logits, targets)
 
-        batch, targets = self._prepare_batch(batch)
+        if not train:
+            self._compute_metrics(logits, targets)
 
-        # get logits
-        logits = self.model(batch)
-        # logits, any
+        return loss, logits
 
-        # compute loss
-        loss = self.model.compute_loss(logits, targets)
-        # loss: dict of losses or loss
+    def training_step(self, batch, batch_idx) -> Any:
+        loss, _ = self.shared_step(batch)
 
         self.log_dict({f"train_{k}": v for k, v in loss.items()}, on_step=True, on_epoch=True, prog_bar=True)
         return loss
@@ -184,17 +167,7 @@ def on_validation_epoch_start(self) -> None:
             metric.reset()
 
     def validation_step(self, batch, batch_idx):
-        batch, targets = self._prepare_batch(batch)
-
-        # get logits
-        logits = self.model(batch)
-        # logits, any
-
-        # compute loss
-        loss = self.model.compute_loss(logits, targets)
-        # loss: dict of losses or loss
-
-        self._compute_metrics(logits, targets)
+        loss, logits = self.shared_step(batch)
 
         self.log_dict({f"val_{k}": v for k, v in loss.items()}, on_step=True, on_epoch=True, prog_bar=True)
         return loss
@@ -208,17 +181,7 @@ def on_test_epoch_start(self) -> None:
             metric.reset()
 
     def test_step(self, batch, batch_idx):
-        batch, targets = self._prepare_batch(batch)
-
-        # get logits
-        logits = self.model(batch)
-        # logits, any
-
-        # compute loss
-        loss = self.model.compute_loss(logits, targets)
-        # loss: dict of losses or loss
-
-        self._compute_metrics(logits, targets)
+        loss, logits = self.shared_step(batch)
 
         self.log_dict({f"test_{k}": v for k, v in loss.items()}, on_step=True, on_epoch=True, prog_bar=True)
         return loss
@@ -233,4 +196,4 @@ def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> A
         return batch
 
     def configure_finetune_callback(self):
-        return [ObjectDetectionFineTuning(train_bn=True)]
+        return [FaceDetectionFineTuning()]
diff --git a/flash_examples/face_detection.py b/flash_examples/face_detection.py
index e8a14f7096..efe00aa4c3 100644
--- a/flash_examples/face_detection.py
+++ b/flash_examples/face_detection.py
@@ -12,30 +12,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import flash
+
+from flash.core.data.utils import download_data
 from flash.core.data.data_module import DataModule
 from flash.core.utilities.imports import _FASTFACE_AVAILABLE
 from flash.image import FaceDetector
-from flash.image.face_detection.data import FaceDetectionPreprocess
+from flash.image.face_detection.data import FaceDetectionPreprocess, FaceDetectionPostProcess
+from flash.image import FaceDetectionData
 
 if _FASTFACE_AVAILABLE:
     import fastface as ff
 else:
-    raise ModuleNotFoundError("Please, pip install -e '.[image]'")
+    raise ModuleNotFoundError("Please, pip install --upgrade 'lightning-flash[image_extras]'")
 
-# 1. Create the DataModule
+# # 1. Create the DataModule
 train_dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="train")
 val_dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="val")
 
-datamodule = DataModule.from_data_source(
-    "fastface", train_data=train_dataset, val_data=val_dataset, preprocess=FaceDetectionPreprocess()
+datamodule = FaceDetectionData.from_datasets(
+    train_dataset=train_dataset, val_dataset=val_dataset, batch_size=2
 )
 
-# 2. Build the task
+# # 2. Build the task
 model = FaceDetector(model="lffd_slim")
 
-# 3. Create the trainer and finetune the model
+# # 3. Create the trainer and finetune the model
 trainer = flash.Trainer(max_epochs=3, limit_train_batches=0.1, limit_val_batches=0.1)
-
 trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
 # 4. Detect faces in a few images!
@@ -46,5 +48,5 @@
 ])
 print(predictions)
 
-# 5. Save the model!
+# # 5. Save the model!
 trainer.save_checkpoint("face_detection_model.pt")

From 13f0c272fa5c80b4befc62ff3c38b8327e4beb30 Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Fri, 24 Sep 2021 07:45:48 -0400
Subject: [PATCH 04/35] pep8

---
 flash/core/data/utils.py               |  2 +-
 flash/image/__init__.py                |  2 +-
 flash/image/face_detection/__init__.py |  2 +-
 flash/image/face_detection/data.py     | 35 ++++++++++----------------
 flash/image/face_detection/model.py    | 16 ++++++------
 flash_examples/face_detection.py       | 23 +++++++----------
 6 files changed, 33 insertions(+), 47 deletions(-)

diff --git a/flash/core/data/utils.py b/flash/core/data/utils.py
index 01b0f2c007..1402414d01 100644
--- a/flash/core/data/utils.py
+++ b/flash/core/data/utils.py
@@ -164,7 +164,7 @@ def extract_tarfile(file_path: str, extract_path: str, mode: str):
                 zip_ref.extractall(path)
     elif local_filename.endswith(".tar.gz") or local_filename.endswith(".tgz"):
         extract_tarfile(local_filename, path, "r:gz")
-    elif local_filename.endswith('.tar.bz2') or local_filename.endswith('.tbz'):
+    elif local_filename.endswith(".tar.bz2") or local_filename.endswith(".tbz"):
         extract_tarfile(local_filename, path, "r:bz2")
 
 
diff --git a/flash/image/__init__.py b/flash/image/__init__.py
index 36257c77ce..788a15ca40 100644
--- a/flash/image/__init__.py
+++ b/flash/image/__init__.py
@@ -6,7 +6,7 @@
 from flash.image.classification.backbones import IMAGE_CLASSIFIER_BACKBONES  # noqa: F401
 from flash.image.detection import ObjectDetectionData, ObjectDetector  # noqa: F401
 from flash.image.embedding import ImageEmbedder  # noqa: F401
-from flash.image.face_detection import FaceDetector, FaceDetectionData # noqa: F401
+from flash.image.face_detection import FaceDetectionData, FaceDetector  # noqa: F401
 from flash.image.instance_segmentation import InstanceSegmentation, InstanceSegmentationData  # noqa: F401
 from flash.image.keypoint_detection import KeypointDetectionData, KeypointDetector  # noqa: F401
 from flash.image.segmentation import (  # noqa: F401
diff --git a/flash/image/face_detection/__init__.py b/flash/image/face_detection/__init__.py
index 7d14bba121..b8a100f085 100644
--- a/flash/image/face_detection/__init__.py
+++ b/flash/image/face_detection/__init__.py
@@ -1,2 +1,2 @@
-from flash.image.face_detection.model import FaceDetector  # noqa: F401
 from flash.image.face_detection.data import FaceDetectionData  # noqa: F401
+from flash.image.face_detection.model import FaceDetector  # noqa: F401
diff --git a/flash/image/face_detection/data.py b/flash/image/face_detection/data.py
index 8d196819e8..7a53ead1e6 100644
--- a/flash/image/face_detection/data.py
+++ b/flash/image/face_detection/data.py
@@ -11,31 +11,23 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Mapping
+from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple
 
 import torch
-import torchvision
 import torch.nn as nn
-
+import torchvision
 from torch.utils.data import Dataset
 
-from flash.core.data.transforms import ApplyToKeys
 from flash.core.data.data_source import DatasetDataSource, DefaultDataKeys, DefaultDataSources
-from flash.core.data.process import Preprocess, Postprocess
-from flash.core.utilities.imports import _TORCHVISION_AVAILABLE, _ICEVISION_AVAILABLE, _FASTFACE_AVAILABLE
+from flash.core.data.process import Postprocess, Preprocess
+from flash.core.data.transforms import ApplyToKeys
+from flash.core.utilities.imports import _FASTFACE_AVAILABLE, _TORCHVISION_AVAILABLE
 from flash.image.data import ImagePathsDataSource
-from flash.core.integrations.icevision.data import IceVisionParserDataSource
-from flash.core.integrations.icevision.transforms import default_transforms
 from flash.image.detection import ObjectDetectionData
 
 if _TORCHVISION_AVAILABLE:
     from torchvision.datasets.folder import default_loader
 
-if _ICEVISION_AVAILABLE:
-    from icevision.parsers import COCOBBoxParser
-else:
-    COCOBBoxParser = object
-
 if _FASTFACE_AVAILABLE:
     import fastface as ff
 
@@ -77,7 +69,7 @@ def load_data(self, data: Dataset, dataset: Any = None) -> Dataset:
                         dict(
                             boxes=targets["target_boxes"],
                             labels=[1 for _ in range(targets["target_boxes"].shape[0])],
-                        )
+                        ),
                     )
                 )
             )
@@ -99,7 +91,6 @@ def load_sample(self, sample: Any, dataset: Optional[Any] = None) -> Mapping[str
 
 
 class FaceDetectionPreprocess(Preprocess):
-
     def __init__(
         self,
         train_transform: Optional[Dict[str, Callable]] = None,
@@ -137,9 +128,9 @@ def default_transforms(self) -> Dict[str, Callable]:
                 ApplyToKeys(
                     DefaultDataKeys.TARGET,
                     nn.Sequential(
-                        ApplyToKeys('boxes', torch.as_tensor),
-                        ApplyToKeys('labels', torch.as_tensor),
-                    )
+                        ApplyToKeys("boxes", torch.as_tensor),
+                        ApplyToKeys("labels", torch.as_tensor),
+                    ),
                 ),
             ),
             "collate": fastface_collate_fn,
@@ -149,11 +140,11 @@ def default_transforms(self) -> Dict[str, Callable]:
 class FaceDetectionPostProcess(Postprocess):
     @staticmethod
     def per_batch_transform(batch: Any) -> Any:
-        scales = batch['scales']
-        paddings = batch['paddings']
+        scales = batch["scales"]
+        paddings = batch["paddings"]
 
-        batch.pop('scales', None)
-        batch.pop('paddings', None)
+        batch.pop("scales", None)
+        batch.pop("paddings", None)
 
         preds = batch[DefaultDataKeys.PREDS]
 
diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index 0679247935..437596b2b7 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -11,21 +11,19 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, List, Mapping, Optional, Sequence, Type, Union, Dict
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Type, Union
 
-import torch
 import pytorch_lightning as pl
-
+import torch
 from torch import nn
 from torch.optim import Optimizer
 
-from flash.core.model import Task
 from flash.core.data.data_source import DefaultDataKeys
-from flash.core.data.process import Postprocess
 from flash.core.data.process import Preprocess, Serializer
-from flash.core.utilities.imports import _FASTFACE_AVAILABLE
 from flash.core.finetuning import FlashBaseFinetuning
-from flash.image.face_detection.data import FaceDetectionPreprocess, FaceDetectionPostProcess
+from flash.core.model import Task
+from flash.core.utilities.imports import _FASTFACE_AVAILABLE
+from flash.image.face_detection.data import FaceDetectionPreprocess
 
 if _FASTFACE_AVAILABLE:
     import fastface as ff
@@ -48,7 +46,9 @@ def serialize(self, sample: Any) -> Dict[str, Any]:
 
 
 class FaceDetector(Task):
-    """The ``FaceDetector`` is a :class:`~flash.Task` for detecting faces in images. For more details, see
+    """The ``FaceDetector`` is a :class:`~flash.Task` for detecting faces in images.
+
+    For more details, see
     :ref:`face_detection`.
     Args:
         model: a string of :attr`_models`. Defaults to 'lffd_slim'.
diff --git a/flash_examples/face_detection.py b/flash_examples/face_detection.py
index efe00aa4c3..4a2d7c7523 100644
--- a/flash_examples/face_detection.py
+++ b/flash_examples/face_detection.py
@@ -12,13 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import flash
-
-from flash.core.data.utils import download_data
-from flash.core.data.data_module import DataModule
 from flash.core.utilities.imports import _FASTFACE_AVAILABLE
-from flash.image import FaceDetector
-from flash.image.face_detection.data import FaceDetectionPreprocess, FaceDetectionPostProcess
-from flash.image import FaceDetectionData
+from flash.image import FaceDetectionData, FaceDetector
 
 if _FASTFACE_AVAILABLE:
     import fastface as ff
@@ -29,9 +24,7 @@
 train_dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="train")
 val_dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="val")
 
-datamodule = FaceDetectionData.from_datasets(
-    train_dataset=train_dataset, val_dataset=val_dataset, batch_size=2
-)
+datamodule = FaceDetectionData.from_datasets(train_dataset=train_dataset, val_dataset=val_dataset, batch_size=2)
 
 # # 2. Build the task
 model = FaceDetector(model="lffd_slim")
@@ -41,11 +34,13 @@
 trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
 # 4. Detect faces in a few images!
-predictions = model.predict([
-    "data/2002/07/19/big/img_18.jpg",
-    "data/2002/07/19/big/img_65.jpg",
-    "data/2002/07/19/big/img_255.jpg",
-])
+predictions = model.predict(
+    [
+        "data/2002/07/19/big/img_18.jpg",
+        "data/2002/07/19/big/img_65.jpg",
+        "data/2002/07/19/big/img_255.jpg",
+    ]
+)
 print(predictions)
 
 # # 5. Save the model!

From 29384d083272aedec8a5a84123fcfa2a232a78f4 Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Fri, 24 Sep 2021 07:53:14 -0400
Subject: [PATCH 05/35] imports

---
 flash/image/face_detection/data.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flash/image/face_detection/data.py b/flash/image/face_detection/data.py
index 7a53ead1e6..605c9f6280 100644
--- a/flash/image/face_detection/data.py
+++ b/flash/image/face_detection/data.py
@@ -15,7 +15,6 @@
 
 import torch
 import torch.nn as nn
-import torchvision
 from torch.utils.data import Dataset
 
 from flash.core.data.data_source import DatasetDataSource, DefaultDataKeys, DefaultDataSources
@@ -26,6 +25,7 @@
 from flash.image.detection import ObjectDetectionData
 
 if _TORCHVISION_AVAILABLE:
+    import torchvision
     from torchvision.datasets.folder import default_loader
 
 if _FASTFACE_AVAILABLE:

From e99adad7424d6e4085a0c6b56aef7ee79a223ce2 Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Fri, 24 Sep 2021 12:21:50 -0400
Subject: [PATCH 06/35] backbones registry

---
 .../face_detection/backbones/__init__.py      |  5 +++
 .../backbones/fastface_backbones.py           | 41 +++++++++++++++++++
 flash/image/face_detection/model.py           |  8 +---
 flash_examples/face_detection.py              |  2 +-
 4 files changed, 49 insertions(+), 7 deletions(-)
 create mode 100644 flash/image/face_detection/backbones/__init__.py
 create mode 100644 flash/image/face_detection/backbones/fastface_backbones.py

diff --git a/flash/image/face_detection/backbones/__init__.py b/flash/image/face_detection/backbones/__init__.py
new file mode 100644
index 0000000000..fdf1185754
--- /dev/null
+++ b/flash/image/face_detection/backbones/__init__.py
@@ -0,0 +1,5 @@
+from flash.core.registry import FlashRegistry  # noqa: F401
+from flash.image.face_detection.backbones.fastface_backbones import register_ff_backbones  # noqa: F401
+
+FACE_DETECTION_BACKBONES = FlashRegistry("face_detection_backbones")
+register_ff_backbones(FACE_DETECTION_BACKBONES)
diff --git a/flash/image/face_detection/backbones/fastface_backbones.py b/flash/image/face_detection/backbones/fastface_backbones.py
new file mode 100644
index 0000000000..5a978cc5d9
--- /dev/null
+++ b/flash/image/face_detection/backbones/fastface_backbones.py
@@ -0,0 +1,41 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from functools import partial
+
+from flash.core.registry import FlashRegistry
+from flash.core.utilities.imports import _FASTFACE_AVAILABLE
+
+if _FASTFACE_AVAILABLE:
+    import fastface as ff
+
+
+model_name = ['lffd_slim', 'lffd_original']
+
+
+def fastface_backbone(model_name, pretrained, **kwargs):
+    if pretrained:
+        model = ff.FaceDetector.from_pretrained(model_name, **kwargs)
+    else:
+        arch, config = model_name.split("_")
+        model = ff.FaceDetector.build(arch, config, **kwargs)
+
+    return model
+
+
+def register_ff_backbones(register: FlashRegistry):
+    backbones = [partial(fastface_backbone, model_name=name) for name in model_name]
+
+    if _FASTFACE_AVAILABLE:
+        for idx, backbone in enumerate(backbones):
+            register(backbone, name=model_name[idx])
diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index 437596b2b7..b6d145d5ef 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -24,6 +24,7 @@
 from flash.core.model import Task
 from flash.core.utilities.imports import _FASTFACE_AVAILABLE
 from flash.image.face_detection.data import FaceDetectionPreprocess
+from flash.image.face_detection.backbones import FACE_DETECTION_BACKBONES
 
 if _FASTFACE_AVAILABLE:
     import fastface as ff
@@ -97,12 +98,7 @@ def get_model(
         pretrained,
         **kwargs,
     ):
-
-        if pretrained:
-            pl_model = ff.FaceDetector.from_pretrained(model_name, **kwargs)
-        else:
-            arch, config = model_name.split("_")
-            pl_model = ff.FaceDetector.build(arch, config, **kwargs)
+        pl_model = FACE_DETECTION_BACKBONES.get(model_name)(pretrained=pretrained, **kwargs)
 
         # get torch.nn.Module
         model = getattr(pl_model, "arch")
diff --git a/flash_examples/face_detection.py b/flash_examples/face_detection.py
index 4a2d7c7523..ec9a00b612 100644
--- a/flash_examples/face_detection.py
+++ b/flash_examples/face_detection.py
@@ -30,7 +30,7 @@
 model = FaceDetector(model="lffd_slim")
 
 # # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=3, limit_train_batches=0.1, limit_val_batches=0.1)
+trainer = flash.Trainer(max_steps=3, limit_train_batches=0.1, limit_val_batches=0.1)
 trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
 # 4. Detect faces in a few images!

From fc7ff44ff77153a03ed18f37de8f0e1c1a9e9e07 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 24 Sep 2021 16:23:15 +0000
Subject: [PATCH 07/35] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 flash/image/face_detection/backbones/fastface_backbones.py | 2 +-
 flash/image/face_detection/model.py                        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/flash/image/face_detection/backbones/fastface_backbones.py b/flash/image/face_detection/backbones/fastface_backbones.py
index 5a978cc5d9..9573ad2429 100644
--- a/flash/image/face_detection/backbones/fastface_backbones.py
+++ b/flash/image/face_detection/backbones/fastface_backbones.py
@@ -20,7 +20,7 @@
     import fastface as ff
 
 
-model_name = ['lffd_slim', 'lffd_original']
+model_name = ["lffd_slim", "lffd_original"]
 
 
 def fastface_backbone(model_name, pretrained, **kwargs):
diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index b6d145d5ef..c5e0dcd3b7 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -23,8 +23,8 @@
 from flash.core.finetuning import FlashBaseFinetuning
 from flash.core.model import Task
 from flash.core.utilities.imports import _FASTFACE_AVAILABLE
-from flash.image.face_detection.data import FaceDetectionPreprocess
 from flash.image.face_detection.backbones import FACE_DETECTION_BACKBONES
+from flash.image.face_detection.data import FaceDetectionPreprocess
 
 if _FASTFACE_AVAILABLE:
     import fastface as ff

From 159bf036b7120c8097d15cf067f47b1b99e8ec3b Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Fri, 24 Sep 2021 12:55:01 -0400
Subject: [PATCH 08/35] tests

---
 flash_examples/face_detection.py         |  3 ++-
 tests/image/face_detection/__init__.py   |  0
 tests/image/face_detection/test_model.py | 32 ++++++++++++++++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)
 create mode 100644 tests/image/face_detection/__init__.py
 create mode 100644 tests/image/face_detection/test_model.py

diff --git a/flash_examples/face_detection.py b/flash_examples/face_detection.py
index ec9a00b612..a1eb002c18 100644
--- a/flash_examples/face_detection.py
+++ b/flash_examples/face_detection.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import torch
 import flash
 from flash.core.utilities.imports import _FASTFACE_AVAILABLE
 from flash.image import FaceDetectionData, FaceDetector
@@ -30,7 +31,7 @@
 model = FaceDetector(model="lffd_slim")
 
 # # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_steps=3, limit_train_batches=0.1, limit_val_batches=0.1)
+trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
 trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
 # 4. Detect faces in a few images!
diff --git a/tests/image/face_detection/__init__.py b/tests/image/face_detection/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/image/face_detection/test_model.py b/tests/image/face_detection/test_model.py
new file mode 100644
index 0000000000..cafc38884a
--- /dev/null
+++ b/tests/image/face_detection/test_model.py
@@ -0,0 +1,32 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import pytest
+
+import flash
+from flash.core.utilities.imports import _FASTFACE_AVAILABLE
+from flash.image import FaceDetectionData, FaceDetector
+
+if  _FASTFACE_AVAILABLE:
+    import fastface as ff
+
+
+@pytest.mark.skipif(not _FASTFACE_AVAILABLE, reason="vissl not installed.")
+def test_fastface_training():
+    dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="val")
+    datamodule = FaceDetectionData.from_datasets(train_dataset=dataset, batch_size=2)
+
+    model = FaceDetector(model="lffd_slim")
+
+    trainer = flash.Trainer(max_steps=2, num_sanity_val_steps=0)
+    trainer.finetune(model, datamodule=datamodule, strategy="freeze")

From 1175b41258849d29b53da5cd46da02b55fb712ff Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Fri, 24 Sep 2021 12:55:52 -0400
Subject: [PATCH 09/35] tests

---
 flash/image/face_detection/backbones/fastface_backbones.py | 2 +-
 flash/image/face_detection/model.py                        | 2 +-
 flash_examples/face_detection.py                           | 1 +
 tests/image/face_detection/test_model.py                   | 2 +-
 4 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/flash/image/face_detection/backbones/fastface_backbones.py b/flash/image/face_detection/backbones/fastface_backbones.py
index 5a978cc5d9..9573ad2429 100644
--- a/flash/image/face_detection/backbones/fastface_backbones.py
+++ b/flash/image/face_detection/backbones/fastface_backbones.py
@@ -20,7 +20,7 @@
     import fastface as ff
 
 
-model_name = ['lffd_slim', 'lffd_original']
+model_name = ["lffd_slim", "lffd_original"]
 
 
 def fastface_backbone(model_name, pretrained, **kwargs):
diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index b6d145d5ef..c5e0dcd3b7 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -23,8 +23,8 @@
 from flash.core.finetuning import FlashBaseFinetuning
 from flash.core.model import Task
 from flash.core.utilities.imports import _FASTFACE_AVAILABLE
-from flash.image.face_detection.data import FaceDetectionPreprocess
 from flash.image.face_detection.backbones import FACE_DETECTION_BACKBONES
+from flash.image.face_detection.data import FaceDetectionPreprocess
 
 if _FASTFACE_AVAILABLE:
     import fastface as ff
diff --git a/flash_examples/face_detection.py b/flash_examples/face_detection.py
index a1eb002c18..6890d17cce 100644
--- a/flash_examples/face_detection.py
+++ b/flash_examples/face_detection.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import torch
+
 import flash
 from flash.core.utilities.imports import _FASTFACE_AVAILABLE
 from flash.image import FaceDetectionData, FaceDetector
diff --git a/tests/image/face_detection/test_model.py b/tests/image/face_detection/test_model.py
index cafc38884a..35de7ad14a 100644
--- a/tests/image/face_detection/test_model.py
+++ b/tests/image/face_detection/test_model.py
@@ -17,7 +17,7 @@
 from flash.core.utilities.imports import _FASTFACE_AVAILABLE
 from flash.image import FaceDetectionData, FaceDetector
 
-if  _FASTFACE_AVAILABLE:
+if _FASTFACE_AVAILABLE:
     import fastface as ff
 
 

From 021b12eff946904eacedc39ecd91d4bc038847e8 Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Fri, 24 Sep 2021 13:16:15 -0400
Subject: [PATCH 10/35] more coverage

---
 tests/image/face_detection/test_model.py | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/tests/image/face_detection/test_model.py b/tests/image/face_detection/test_model.py
index 35de7ad14a..3e102bf53f 100644
--- a/tests/image/face_detection/test_model.py
+++ b/tests/image/face_detection/test_model.py
@@ -14,14 +14,21 @@
 import pytest
 
 import flash
+from flash.core.registry import FlashRegistry
 from flash.core.utilities.imports import _FASTFACE_AVAILABLE
 from flash.image import FaceDetectionData, FaceDetector
 
 if _FASTFACE_AVAILABLE:
     import fastface as ff
 
+    from fastface.arch.lffd import LFFD
+    from flash.image.face_detection.backbones import FACE_DETECTION_BACKBONES
+else:
+    FACE_DETECTION_BACKBONES = FlashRegistry("face_detection_backbones")
+    LFFD = object
 
-@pytest.mark.skipif(not _FASTFACE_AVAILABLE, reason="vissl not installed.")
+
+@pytest.mark.skipif(not _FASTFACE_AVAILABLE, reason="fastface not installed.")
 def test_fastface_training():
     dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="val")
     datamodule = FaceDetectionData.from_datasets(train_dataset=dataset, batch_size=2)
@@ -30,3 +37,13 @@ def test_fastface_training():
 
     trainer = flash.Trainer(max_steps=2, num_sanity_val_steps=0)
     trainer.finetune(model, datamodule=datamodule, strategy="freeze")
+
+
+@pytest.mark.skipif(not _FASTFACE_AVAILABLE, reason="fastface not installed.")
+def test_fastface_backbones_registry():
+    backbones = FACE_DETECTION_BACKBONES.available_keys()
+    assert 'lffd_slim' in backbones
+    assert 'lffd_original' in backbones
+
+    backbone = FACE_DETECTION_BACKBONES.get('lffd_original')
+    assert isinstance(backbone, LFFD)

From fdff7744f38489241cbffd1735ac12571685e678 Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Fri, 24 Sep 2021 13:48:02 -0400
Subject: [PATCH 11/35] final

---
 .../face_detection/backbones/fastface_backbones.py   | 12 +++++++-----
 flash/image/face_detection/model.py                  |  5 +----
 tests/image/face_detection/test_model.py             |  8 ++++----
 3 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/flash/image/face_detection/backbones/fastface_backbones.py b/flash/image/face_detection/backbones/fastface_backbones.py
index 9573ad2429..de32a5c169 100644
--- a/flash/image/face_detection/backbones/fastface_backbones.py
+++ b/flash/image/face_detection/backbones/fastface_backbones.py
@@ -25,17 +25,19 @@
 
 def fastface_backbone(model_name, pretrained, **kwargs):
     if pretrained:
-        model = ff.FaceDetector.from_pretrained(model_name, **kwargs)
+        pl_model = ff.FaceDetector.from_pretrained(model_name, **kwargs)
     else:
         arch, config = model_name.split("_")
-        model = ff.FaceDetector.build(arch, config, **kwargs)
+        pl_model = ff.FaceDetector.build(arch, config, **kwargs)
 
-    return model
+    backbone = getattr(pl_model, "arch")
 
+    return backbone, pl_model
 
-def register_ff_backbones(register: FlashRegistry):
-    backbones = [partial(fastface_backbone, model_name=name) for name in model_name]
 
+def register_ff_backbones(register: FlashRegistry):
     if _FASTFACE_AVAILABLE:
+        backbones = [partial(fastface_backbone, model_name=name) for name in model_name]
+
         for idx, backbone in enumerate(backbones):
             register(backbone, name=model_name[idx])
diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index c5e0dcd3b7..f6bc41b115 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -98,10 +98,7 @@ def get_model(
         pretrained,
         **kwargs,
     ):
-        pl_model = FACE_DETECTION_BACKBONES.get(model_name)(pretrained=pretrained, **kwargs)
-
-        # get torch.nn.Module
-        model = getattr(pl_model, "arch")
+        model, pl_model = FACE_DETECTION_BACKBONES.get(model_name)(pretrained=pretrained, **kwargs)
 
         # set preprocess params
         model.register_buffer("normalizer", getattr(pl_model, "normalizer"))
diff --git a/tests/image/face_detection/test_model.py b/tests/image/face_detection/test_model.py
index 3e102bf53f..abdec7d8ae 100644
--- a/tests/image/face_detection/test_model.py
+++ b/tests/image/face_detection/test_model.py
@@ -20,8 +20,8 @@
 
 if _FASTFACE_AVAILABLE:
     import fastface as ff
-
     from fastface.arch.lffd import LFFD
+
     from flash.image.face_detection.backbones import FACE_DETECTION_BACKBONES
 else:
     FACE_DETECTION_BACKBONES = FlashRegistry("face_detection_backbones")
@@ -42,8 +42,8 @@ def test_fastface_training():
 @pytest.mark.skipif(not _FASTFACE_AVAILABLE, reason="fastface not installed.")
 def test_fastface_backbones_registry():
     backbones = FACE_DETECTION_BACKBONES.available_keys()
-    assert 'lffd_slim' in backbones
-    assert 'lffd_original' in backbones
+    assert "lffd_slim" in backbones
+    assert "lffd_original" in backbones
 
-    backbone = FACE_DETECTION_BACKBONES.get('lffd_original')
+    backbone, _ = FACE_DETECTION_BACKBONES.get("lffd_original")(pretrained=True)
     assert isinstance(backbone, LFFD)

From 97795672eaef309c6b8f48bf80163a122e6b547f Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Mon, 27 Sep 2021 10:54:42 -0400
Subject: [PATCH 12/35] .

---
 .../face_detection/backbones/fastface_backbones.py   | 12 ++++++------
 flash/image/face_detection/data.py                   |  2 ++
 flash/image/face_detection/model.py                  | 11 +----------
 flash_examples/face_detection.py                     | 10 +++++-----
 4 files changed, 14 insertions(+), 21 deletions(-)

diff --git a/flash/image/face_detection/backbones/fastface_backbones.py b/flash/image/face_detection/backbones/fastface_backbones.py
index de32a5c169..8e59662120 100644
--- a/flash/image/face_detection/backbones/fastface_backbones.py
+++ b/flash/image/face_detection/backbones/fastface_backbones.py
@@ -14,13 +14,13 @@
 from functools import partial
 
 from flash.core.registry import FlashRegistry
-from flash.core.utilities.imports import _FASTFACE_AVAILABLE
+from flash.core.utilities.imports import _FASTFACE_AVAILABLE, requires
 
 if _FASTFACE_AVAILABLE:
     import fastface as ff
 
 
-model_name = ["lffd_slim", "lffd_original"]
+_MODEL_NAMES = ff.list_pretrained_models()
 
 
 def fastface_backbone(model_name, pretrained, **kwargs):
@@ -35,9 +35,9 @@ def fastface_backbone(model_name, pretrained, **kwargs):
     return backbone, pl_model
 
 
+@requires("fastface")
 def register_ff_backbones(register: FlashRegistry):
-    if _FASTFACE_AVAILABLE:
-        backbones = [partial(fastface_backbone, model_name=name) for name in model_name]
+    backbones = [partial(fastface_backbone, model_name=name) for name in _MODEL_NAMES]
 
-        for idx, backbone in enumerate(backbones):
-            register(backbone, name=model_name[idx])
+    for idx, backbone in enumerate(backbones):
+        register(backbone, name=_MODEL_NAMES[idx])
diff --git a/flash/image/face_detection/data.py b/flash/image/face_detection/data.py
index 605c9f6280..0c1a96e72c 100644
--- a/flash/image/face_detection/data.py
+++ b/flash/image/face_detection/data.py
@@ -33,6 +33,8 @@
 
 
 def fastface_collate_fn(samples: Sequence[Dict[str, Any]]) -> Dict[str, Sequence[Any]]:
+    """
+    """
     samples = {key: [sample[key] for sample in samples] for key in samples[0]}
 
     images, scales, paddings = ff.utils.preprocess.prepare_batch(
diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index f6bc41b115..b280fbd933 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -144,8 +144,7 @@ def shared_step(self, batch, train=False) -> Any:
         logits = self.model(images)
         loss = self.model.compute_loss(logits, targets)
 
-        if not train:
-            self._compute_metrics(logits, targets)
+        self._compute_metrics(logits, targets)
 
         return loss, logits
 
@@ -155,10 +154,6 @@ def training_step(self, batch, batch_idx) -> Any:
         self.log_dict({f"train_{k}": v for k, v in loss.items()}, on_step=True, on_epoch=True, prog_bar=True)
         return loss
 
-    def on_validation_epoch_start(self) -> None:
-        for metric in self.val_metrics.values():
-            metric.reset()
-
     def validation_step(self, batch, batch_idx):
         loss, logits = self.shared_step(batch)
 
@@ -169,10 +164,6 @@ def validation_epoch_end(self, outputs) -> None:
         metric_results = {name: metric.compute() for name, metric in self.val_metrics.items()}
         self.log_dict({f"val_{k}": v for k, v in metric_results.items()}, on_epoch=True)
 
-    def on_test_epoch_start(self) -> None:
-        for metric in self.val_metrics.values():
-            metric.reset()
-
     def test_step(self, batch, batch_idx):
         loss, logits = self.shared_step(batch)
 
diff --git a/flash_examples/face_detection.py b/flash_examples/face_detection.py
index 6890d17cce..99f82ba9ba 100644
--- a/flash_examples/face_detection.py
+++ b/flash_examples/face_detection.py
@@ -14,13 +14,13 @@
 import torch
 
 import flash
-from flash.core.utilities.imports import _FASTFACE_AVAILABLE
+from flash.core.utilities.imports import _FASTFACE_AVAILABLE, example_requires
 from flash.image import FaceDetectionData, FaceDetector
 
-if _FASTFACE_AVAILABLE:
-    import fastface as ff
-else:
-    raise ModuleNotFoundError("Please, pip install --upgrade 'lightning-flash[image_extras]'")
+example_requires("fastface")
+
+import fastface as ff  # noqa: E402
+
 
 # # 1. Create the DataModule
 train_dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="train")

From b4dbb96b2ec6c13a27ae1259e6e99b827642a176 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 27 Sep 2021 14:55:50 +0000
Subject: [PATCH 13/35] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 flash/image/face_detection/data.py | 3 +--
 flash_examples/face_detection.py   | 1 -
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/flash/image/face_detection/data.py b/flash/image/face_detection/data.py
index 0c1a96e72c..e640612dd4 100644
--- a/flash/image/face_detection/data.py
+++ b/flash/image/face_detection/data.py
@@ -33,8 +33,7 @@
 
 
 def fastface_collate_fn(samples: Sequence[Dict[str, Any]]) -> Dict[str, Sequence[Any]]:
-    """
-    """
+    """"""
     samples = {key: [sample[key] for sample in samples] for key in samples[0]}
 
     images, scales, paddings = ff.utils.preprocess.prepare_batch(
diff --git a/flash_examples/face_detection.py b/flash_examples/face_detection.py
index 99f82ba9ba..82f6b9dbe4 100644
--- a/flash_examples/face_detection.py
+++ b/flash_examples/face_detection.py
@@ -21,7 +21,6 @@
 
 import fastface as ff  # noqa: E402
 
-
 # # 1. Create the DataModule
 train_dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="train")
 val_dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="val")

From b7af234164e4d0c05d00e92b7bada38ee3ef54d7 Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Mon, 27 Sep 2021 11:00:20 -0400
Subject: [PATCH 14/35] .

---
 flash_examples/face_detection.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/flash_examples/face_detection.py b/flash_examples/face_detection.py
index 99f82ba9ba..33f42dc878 100644
--- a/flash_examples/face_detection.py
+++ b/flash_examples/face_detection.py
@@ -17,10 +17,10 @@
 from flash.core.utilities.imports import _FASTFACE_AVAILABLE, example_requires
 from flash.image import FaceDetectionData, FaceDetector
 
-example_requires("fastface")
-
-import fastface as ff  # noqa: E402
-
+if _FASTFACE_AVAILABLE:
+    import fastface as ff
+else:
+    example_requires("fastface")
 
 # # 1. Create the DataModule
 train_dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="train")

From 70514f0dc2da3e5fc8451218ea4d85cd126653d1 Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Mon, 27 Sep 2021 11:06:00 -0400
Subject: [PATCH 15/35] .

---
 .../image/face_detection/backbones/fastface_backbones.py  | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/flash/image/face_detection/backbones/fastface_backbones.py b/flash/image/face_detection/backbones/fastface_backbones.py
index 8e59662120..52e5418f2c 100644
--- a/flash/image/face_detection/backbones/fastface_backbones.py
+++ b/flash/image/face_detection/backbones/fastface_backbones.py
@@ -35,9 +35,9 @@ def fastface_backbone(model_name, pretrained, **kwargs):
     return backbone, pl_model
 
 
-@requires("fastface")
 def register_ff_backbones(register: FlashRegistry):
-    backbones = [partial(fastface_backbone, model_name=name) for name in _MODEL_NAMES]
+    if _FASTFACE_AVAILABLE:
+        backbones = [partial(fastface_backbone, model_name=name) for name in _MODEL_NAMES]
 
-    for idx, backbone in enumerate(backbones):
-        register(backbone, name=_MODEL_NAMES[idx])
+        for idx, backbone in enumerate(backbones):
+            register(backbone, name=_MODEL_NAMES[idx])

From fceeb2dec46eebbe14917ce197ab2beff93da27d Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Mon, 27 Sep 2021 11:16:23 -0400
Subject: [PATCH 16/35] .

---
 flash/image/face_detection/backbones/fastface_backbones.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/flash/image/face_detection/backbones/fastface_backbones.py b/flash/image/face_detection/backbones/fastface_backbones.py
index 52e5418f2c..dde12b35a9 100644
--- a/flash/image/face_detection/backbones/fastface_backbones.py
+++ b/flash/image/face_detection/backbones/fastface_backbones.py
@@ -19,8 +19,9 @@
 if _FASTFACE_AVAILABLE:
     import fastface as ff
 
-
-_MODEL_NAMES = ff.list_pretrained_models()
+    _MODEL_NAMES = ff.list_pretrained_models()
+else:
+    _MODEL_NAMES = []
 
 
 def fastface_backbone(model_name, pretrained, **kwargs):

From 24fcdbdccc3d8ce00df8cb47f1da0fae96b34749 Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Mon, 27 Sep 2021 11:18:28 -0400
Subject: [PATCH 17/35] .

---
 flash/image/face_detection/backbones/fastface_backbones.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flash/image/face_detection/backbones/fastface_backbones.py b/flash/image/face_detection/backbones/fastface_backbones.py
index dde12b35a9..57062ab130 100644
--- a/flash/image/face_detection/backbones/fastface_backbones.py
+++ b/flash/image/face_detection/backbones/fastface_backbones.py
@@ -14,7 +14,7 @@
 from functools import partial
 
 from flash.core.registry import FlashRegistry
-from flash.core.utilities.imports import _FASTFACE_AVAILABLE, requires
+from flash.core.utilities.imports import _FASTFACE_AVAILABLE
 
 if _FASTFACE_AVAILABLE:
     import fastface as ff

From 3e2eeb009a58566eb58043196b378a33d321cf0a Mon Sep 17 00:00:00 2001
From: Ananya Harsh Jha <ananya@pytorchlightning.ai>
Date: Mon, 27 Sep 2021 11:22:12 -0400
Subject: [PATCH 18/35] Update flash/image/face_detection/model.py

Co-authored-by: Sean Naren <sean@grid.ai>
---
 flash/image/face_detection/model.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index b280fbd933..2a7a8ff517 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -42,8 +42,7 @@ class DetectionLabels(Serializer):
     """A :class:`.Serializer` which extracts predictions from sample dict."""
 
     def serialize(self, sample: Any) -> Dict[str, Any]:
-        sample = sample[DefaultDataKeys.PREDS] if isinstance(sample, Dict) else sample
-        return sample
+        return sample[DefaultDataKeys.PREDS] if isinstance(sample, Dict) else sample
 
 
 class FaceDetector(Task):

From f1bff2b8f4730814072ff743e1da3a3d89ffe3ab Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Mon, 27 Sep 2021 11:33:40 -0400
Subject: [PATCH 19/35] comments

---
 flash/image/face_detection/data.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/flash/image/face_detection/data.py b/flash/image/face_detection/data.py
index e640612dd4..364cf9dbd4 100644
--- a/flash/image/face_detection/data.py
+++ b/flash/image/face_detection/data.py
@@ -33,7 +33,9 @@
 
 
 def fastface_collate_fn(samples: Sequence[Dict[str, Any]]) -> Dict[str, Sequence[Any]]:
-    """"""
+    """Collate function from fastface.
+    Organizes individual elements in a batch, calls prepare_batch from fastface and prepares the targets.
+    """
     samples = {key: [sample[key] for sample in samples] for key in samples[0]}
 
     images, scales, paddings = ff.utils.preprocess.prepare_batch(
@@ -60,6 +62,8 @@ def fastface_collate_fn(samples: Sequence[Dict[str, Any]]) -> Dict[str, Sequence
 
 
 class FastFaceDataSource(DatasetDataSource):
+    """Logic for loading from FDDBDataset.
+    """
     def load_data(self, data: Dataset, dataset: Any = None) -> Dataset:
         new_data = []
         for img_file_path, targets in zip(data.ids, data.targets):
@@ -92,6 +96,8 @@ def load_sample(self, sample: Any, dataset: Optional[Any] = None) -> Mapping[str
 
 
 class FaceDetectionPreprocess(Preprocess):
+    """Applies default transform and collate_fn for fastface on FastFaceDataSource
+    """
     def __init__(
         self,
         train_transform: Optional[Dict[str, Callable]] = None,
@@ -139,6 +145,8 @@ def default_transforms(self) -> Dict[str, Callable]:
 
 
 class FaceDetectionPostProcess(Postprocess):
+    """Generates preds from model output.
+    """
     @staticmethod
     def per_batch_transform(batch: Any) -> Any:
         scales = batch["scales"]

From 5ef68120d93f0bea1e62446e643e505dfa3428ef Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 27 Sep 2021 15:34:32 +0000
Subject: [PATCH 20/35] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 flash/image/face_detection/data.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/flash/image/face_detection/data.py b/flash/image/face_detection/data.py
index 364cf9dbd4..5f20d1d875 100644
--- a/flash/image/face_detection/data.py
+++ b/flash/image/face_detection/data.py
@@ -34,6 +34,7 @@
 
 def fastface_collate_fn(samples: Sequence[Dict[str, Any]]) -> Dict[str, Sequence[Any]]:
     """Collate function from fastface.
+
     Organizes individual elements in a batch, calls prepare_batch from fastface and prepares the targets.
     """
     samples = {key: [sample[key] for sample in samples] for key in samples[0]}
@@ -62,8 +63,8 @@ def fastface_collate_fn(samples: Sequence[Dict[str, Any]]) -> Dict[str, Sequence
 
 
 class FastFaceDataSource(DatasetDataSource):
-    """Logic for loading from FDDBDataset.
-    """
+    """Logic for loading from FDDBDataset."""
+
     def load_data(self, data: Dataset, dataset: Any = None) -> Dataset:
         new_data = []
         for img_file_path, targets in zip(data.ids, data.targets):
@@ -96,8 +97,8 @@ def load_sample(self, sample: Any, dataset: Optional[Any] = None) -> Mapping[str
 
 
 class FaceDetectionPreprocess(Preprocess):
-    """Applies default transform and collate_fn for fastface on FastFaceDataSource
-    """
+    """Applies default transform and collate_fn for fastface on FastFaceDataSource."""
+
     def __init__(
         self,
         train_transform: Optional[Dict[str, Callable]] = None,
@@ -145,8 +146,8 @@ def default_transforms(self) -> Dict[str, Callable]:
 
 
 class FaceDetectionPostProcess(Postprocess):
-    """Generates preds from model output.
-    """
+    """Generates preds from model output."""
+
     @staticmethod
     def per_batch_transform(batch: Any) -> Any:
         scales = batch["scales"]

From 087a3176cc9452968007ee75c5b67e22a0e66d7d Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Mon, 27 Sep 2021 12:56:59 -0400
Subject: [PATCH 21/35] .

---
 flash/image/face_detection/data.py  | 13 +++++++------
 flash/image/face_detection/model.py | 18 ++++++++++--------
 flash_examples/face_detection.py    |  8 +++-----
 tests/core/test_utils.py            | 10 +++++++++-
 4 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/flash/image/face_detection/data.py b/flash/image/face_detection/data.py
index 364cf9dbd4..5f20d1d875 100644
--- a/flash/image/face_detection/data.py
+++ b/flash/image/face_detection/data.py
@@ -34,6 +34,7 @@
 
 def fastface_collate_fn(samples: Sequence[Dict[str, Any]]) -> Dict[str, Sequence[Any]]:
     """Collate function from fastface.
+
     Organizes individual elements in a batch, calls prepare_batch from fastface and prepares the targets.
     """
     samples = {key: [sample[key] for sample in samples] for key in samples[0]}
@@ -62,8 +63,8 @@ def fastface_collate_fn(samples: Sequence[Dict[str, Any]]) -> Dict[str, Sequence
 
 
 class FastFaceDataSource(DatasetDataSource):
-    """Logic for loading from FDDBDataset.
-    """
+    """Logic for loading from FDDBDataset."""
+
     def load_data(self, data: Dataset, dataset: Any = None) -> Dataset:
         new_data = []
         for img_file_path, targets in zip(data.ids, data.targets):
@@ -96,8 +97,8 @@ def load_sample(self, sample: Any, dataset: Optional[Any] = None) -> Mapping[str
 
 
 class FaceDetectionPreprocess(Preprocess):
-    """Applies default transform and collate_fn for fastface on FastFaceDataSource
-    """
+    """Applies default transform and collate_fn for fastface on FastFaceDataSource."""
+
     def __init__(
         self,
         train_transform: Optional[Dict[str, Callable]] = None,
@@ -145,8 +146,8 @@ def default_transforms(self) -> Dict[str, Callable]:
 
 
 class FaceDetectionPostProcess(Postprocess):
-    """Generates preds from model output.
-    """
+    """Generates preds from model output."""
+
     @staticmethod
     def per_batch_transform(batch: Any) -> Any:
         scales = batch["scales"]
diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index 2a7a8ff517..5ae9337200 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -84,7 +84,7 @@ def __init__(
         super().__init__(
             model=model,
             loss_fn=loss,
-            metrics=metrics or {"AP": ff.metric.AveragePrecision()},
+            metrics=metrics or {"AP": ff.metric.AveragePrecision()},  # TODO: replace with torch metrics MAP
             learning_rate=learning_rate,
             optimizer=optimizer,
             serializer=serializer or DetectionLabels(),
@@ -105,6 +105,8 @@ def get_model(
         model.register_buffer("std", getattr(pl_model, "std"))
 
         # set postprocess function
+        # this is called from FaceDetector lightning module form fastface itself
+        # https://github.com/borhanMorphy/fastface/blob/master/fastface/module.py#L200
         setattr(model, "_postprocess", getattr(pl_model, "_postprocess"))
 
         return model
@@ -125,11 +127,11 @@ def _prepare_batch(self, batch):
         return batch
 
     def _compute_metrics(self, logits, targets):
-        preds = self.model.logits_to_preds(logits)
         # preds: torch.Tensor(B, N, 5)
+        preds = self.model.logits_to_preds(logits)
 
-        preds = self.model._postprocess(preds)
         # preds: torch.Tensor(N, 6) as x1,y1,x2,y2,score,batch_idx
+        preds = self.model._postprocess(preds)
 
         target_boxes = [target["target_boxes"] for target in targets]
         pred_boxes = [preds[preds[:, 5] == batch_idx, :5] for batch_idx in range(len(targets))]
@@ -137,7 +139,7 @@ def _compute_metrics(self, logits, targets):
         for metric in self.val_metrics.values():
             metric.update(pred_boxes, target_boxes)
 
-    def shared_step(self, batch, train=False) -> Any:
+    def __shared_step(self, batch, train=False) -> Any:
         images, targets = batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET]
         images = self._prepare_batch(images)
         logits = self.model(images)
@@ -145,16 +147,16 @@ def shared_step(self, batch, train=False) -> Any:
 
         self._compute_metrics(logits, targets)
 
-        return loss, logits
+        return loss
 
     def training_step(self, batch, batch_idx) -> Any:
-        loss, _ = self.shared_step(batch)
+        loss, _ = self.__shared_step(batch)
 
         self.log_dict({f"train_{k}": v for k, v in loss.items()}, on_step=True, on_epoch=True, prog_bar=True)
         return loss
 
     def validation_step(self, batch, batch_idx):
-        loss, logits = self.shared_step(batch)
+        loss = self.__shared_step(batch)
 
         self.log_dict({f"val_{k}": v for k, v in loss.items()}, on_step=True, on_epoch=True, prog_bar=True)
         return loss
@@ -164,7 +166,7 @@ def validation_epoch_end(self, outputs) -> None:
         self.log_dict({f"val_{k}": v for k, v in metric_results.items()}, on_epoch=True)
 
     def test_step(self, batch, batch_idx):
-        loss, logits = self.shared_step(batch)
+        loss = self.__shared_step(batch)
 
         self.log_dict({f"test_{k}": v for k, v in loss.items()}, on_step=True, on_epoch=True, prog_bar=True)
         return loss
diff --git a/flash_examples/face_detection.py b/flash_examples/face_detection.py
index 33f42dc878..7c5f35b4ab 100644
--- a/flash_examples/face_detection.py
+++ b/flash_examples/face_detection.py
@@ -14,13 +14,11 @@
 import torch
 
 import flash
-from flash.core.utilities.imports import _FASTFACE_AVAILABLE, example_requires
+from flash.core.utilities.imports import example_requires
 from flash.image import FaceDetectionData, FaceDetector
 
-if _FASTFACE_AVAILABLE:
-    import fastface as ff
-else:
-    example_requires("fastface")
+example_requires("fastface")
+import fastface as ff
 
 # # 1. Create the DataModule
 train_dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="train")
diff --git a/tests/core/test_utils.py b/tests/core/test_utils.py
index 49d24bf7ab..8b2f569553 100644
--- a/tests/core/test_utils.py
+++ b/tests/core/test_utils.py
@@ -53,4 +53,12 @@ def test_get_callable_dict():
 def test_download_data(tmpdir):
     path = os.path.join(tmpdir, "data")
     download_data("https://pl-flash-data.s3.amazonaws.com/titanic.zip", path)
-    assert set(os.listdir(path)) == {"titanic", "titanic.zip"}
+    assert "titanic" in set(os.listdir(path))
+    assert "titanic.zip" in set(os.listdir(path))
+
+
+def test_download_data_tar(tmpdir):
+    path = os.path.join(tmpdir, "data")
+    download_data("https://pl-flash-data.s3.amazonaws.com/titanic.tar.gz", path)
+    assert "titanic" in set(os.listdir(path))
+    assert "titanic.tar.gz" in set(os.listdir(path))

From d570cb3328c0a0420443c3b58455d93b1e38ed4f Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Mon, 27 Sep 2021 13:12:09 -0400
Subject: [PATCH 22/35] .

---
 flash_examples/face_detection.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/flash_examples/face_detection.py b/flash_examples/face_detection.py
index 7c5f35b4ab..1457e290ad 100644
--- a/flash_examples/face_detection.py
+++ b/flash_examples/face_detection.py
@@ -18,7 +18,8 @@
 from flash.image import FaceDetectionData, FaceDetector
 
 example_requires("fastface")
-import fastface as ff
+
+import fastface as ff  # noqa: F401
 
 # # 1. Create the DataModule
 train_dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="train")

From bbda29486e4963ac5859a85595d693c16db5d26b Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Mon, 27 Sep 2021 15:45:11 -0400
Subject: [PATCH 23/35] .

---
 tests/core/test_utils.py | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/tests/core/test_utils.py b/tests/core/test_utils.py
index 8b2f569553..99a2cd5e18 100644
--- a/tests/core/test_utils.py
+++ b/tests/core/test_utils.py
@@ -55,10 +55,3 @@ def test_download_data(tmpdir):
     download_data("https://pl-flash-data.s3.amazonaws.com/titanic.zip", path)
     assert "titanic" in set(os.listdir(path))
     assert "titanic.zip" in set(os.listdir(path))
-
-
-def test_download_data_tar(tmpdir):
-    path = os.path.join(tmpdir, "data")
-    download_data("https://pl-flash-data.s3.amazonaws.com/titanic.tar.gz", path)
-    assert "titanic" in set(os.listdir(path))
-    assert "titanic.tar.gz" in set(os.listdir(path))

From 5db8acacb681882388547dee6644a9ce00edecde Mon Sep 17 00:00:00 2001
From: borhanMorphy <borhano.f.42@gmail.com>
Date: Wed, 29 Sep 2021 01:48:29 +0300
Subject: [PATCH 24/35] added comments to clearfy some steps in the face
 detection task

---
 flash/image/face_detection/data.py  | 1 +
 flash/image/face_detection/model.py | 3 +++
 2 files changed, 4 insertions(+)

diff --git a/flash/image/face_detection/data.py b/flash/image/face_detection/data.py
index 7a53ead1e6..f313f3998e 100644
--- a/flash/image/face_detection/data.py
+++ b/flash/image/face_detection/data.py
@@ -68,6 +68,7 @@ def load_data(self, data: Dataset, dataset: Any = None) -> Dataset:
                         img_file_path,
                         dict(
                             boxes=targets["target_boxes"],
+                            # label `1` indicates positive sample
                             labels=[1 for _ in range(targets["target_boxes"].shape[0])],
                         ),
                     )
diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index 437596b2b7..7684e1bc7e 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -104,14 +104,17 @@ def get_model(
             arch, config = model_name.split("_")
             pl_model = ff.FaceDetector.build(arch, config, **kwargs)
 
+        # following steps are required since `get_model` needsreturns `torch.nn.Module`
         # get torch.nn.Module
         model = getattr(pl_model, "arch")
 
+        # moving some required parameters from `fastface.FaceDetector` to `torch.nn.Module`
         # set preprocess params
         model.register_buffer("normalizer", getattr(pl_model, "normalizer"))
         model.register_buffer("mean", getattr(pl_model, "mean"))
         model.register_buffer("std", getattr(pl_model, "std"))
 
+        # copy pasting `_postprocess` function from `fastface.FaceDetector` to `torch.nn.Module`
         # set postprocess function
         setattr(model, "_postprocess", getattr(pl_model, "_postprocess"))
 

From e2ca442af3b7cfc3d81d4ac37115fd09b812f165 Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Tue, 28 Sep 2021 18:59:19 -0400
Subject: [PATCH 25/35] imports

---
 flash_examples/face_detection.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/flash_examples/face_detection.py b/flash_examples/face_detection.py
index 1457e290ad..9762cadb01 100644
--- a/flash_examples/face_detection.py
+++ b/flash_examples/face_detection.py
@@ -18,8 +18,7 @@
 from flash.image import FaceDetectionData, FaceDetector
 
 example_requires("fastface")
-
-import fastface as ff  # noqa: F401
+import fastface as ff  # noqa: E402
 
 # # 1. Create the DataModule
 train_dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="train")

From 4deaf9c88691616f2a55c5e5258f9104ecda8235 Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Tue, 28 Sep 2021 20:47:24 -0400
Subject: [PATCH 26/35] .

---
 flash/image/face_detection/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index b57ed53382..748a44292b 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -153,7 +153,7 @@ def __shared_step(self, batch, train=False) -> Any:
         return loss
 
     def training_step(self, batch, batch_idx) -> Any:
-        loss, _ = self.__shared_step(batch)
+        loss = self.__shared_step(batch)
 
         self.log_dict({f"train_{k}": v for k, v in loss.items()}, on_step=True, on_epoch=True, prog_bar=True)
         return loss

From 2b7c3260751079212c018b1e397fc4d8027c272e Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Tue, 28 Sep 2021 21:01:34 -0400
Subject: [PATCH 27/35] .

---
 tests/core/test_utils.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/core/test_utils.py b/tests/core/test_utils.py
index 99a2cd5e18..ff68e99292 100644
--- a/tests/core/test_utils.py
+++ b/tests/core/test_utils.py
@@ -55,3 +55,9 @@ def test_download_data(tmpdir):
     download_data("https://pl-flash-data.s3.amazonaws.com/titanic.zip", path)
     assert "titanic" in set(os.listdir(path))
     assert "titanic.zip" in set(os.listdir(path))
+
+def test_download_targz_data(tmpdir):
+    path = os.path.join(tmpdir, "data")
+    download_data("https://pl-flash-data.s3.amazonaws.com/titanic.tar.gz", path)
+    assert "titanic" in set(os.listdir(path))
+    assert "titanic.tar.gz" in set(os.listdir(path))

From 144e063873df817fd57fe9396cc42d775706d2d7 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 29 Sep 2021 01:02:15 +0000
Subject: [PATCH 28/35] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/core/test_utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/core/test_utils.py b/tests/core/test_utils.py
index ff68e99292..fdf941744a 100644
--- a/tests/core/test_utils.py
+++ b/tests/core/test_utils.py
@@ -56,6 +56,7 @@ def test_download_data(tmpdir):
     assert "titanic" in set(os.listdir(path))
     assert "titanic.zip" in set(os.listdir(path))
 
+
 def test_download_targz_data(tmpdir):
     path = os.path.join(tmpdir, "data")
     download_data("https://pl-flash-data.s3.amazonaws.com/titanic.tar.gz", path)

From 2085764fd0be0f589d730752cf3d460b67386630 Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Tue, 28 Sep 2021 21:29:17 -0400
Subject: [PATCH 29/35] .

---
 tests/core/test_utils.py | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/tests/core/test_utils.py b/tests/core/test_utils.py
index ff68e99292..184cce8da3 100644
--- a/tests/core/test_utils.py
+++ b/tests/core/test_utils.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+import pytest
 
 from flash.core.data.utils import download_data
 from flash.core.utilities.apply_func import get_callable_dict, get_callable_name
@@ -50,14 +51,10 @@ def test_get_callable_dict():
     assert d["two"] == b
 
 
-def test_download_data(tmpdir):
+@pytest.mark.parametrize("file", ["titanic.zip", "titanic.tar.gz", "titanic.tar.bz2"])
+def test_download_data(tmpdir, file):
+    download_path = "https://pl-flash-data.s3.amazonaws.com/"
     path = os.path.join(tmpdir, "data")
-    download_data("https://pl-flash-data.s3.amazonaws.com/titanic.zip", path)
+    download_data(download_path + file, path)
     assert "titanic" in set(os.listdir(path))
-    assert "titanic.zip" in set(os.listdir(path))
-
-def test_download_targz_data(tmpdir):
-    path = os.path.join(tmpdir, "data")
-    download_data("https://pl-flash-data.s3.amazonaws.com/titanic.tar.gz", path)
-    assert "titanic" in set(os.listdir(path))
-    assert "titanic.tar.gz" in set(os.listdir(path))
+    assert file in set(os.listdir(path))

From 0e4c86225b11e3cf7f3036dd028b70b04ae99bbe Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Tue, 28 Sep 2021 21:30:42 -0400
Subject: [PATCH 30/35] .

---
 tests/core/test_utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/core/test_utils.py b/tests/core/test_utils.py
index 184cce8da3..61666b70ca 100644
--- a/tests/core/test_utils.py
+++ b/tests/core/test_utils.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+
 import pytest
 
 from flash.core.data.utils import download_data

From 6a9a9a1cdf5bba99a4a4baf479d6f2a68e04ae99 Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Thu, 30 Sep 2021 09:34:59 -0400
Subject: [PATCH 31/35] tests

---
 flash/core/data/utils.py                                 | 2 +-
 .../image/face_detection/backbones/fastface_backbones.py | 2 +-
 flash/image/face_detection/model.py                      | 6 +++---
 tests/image/face_detection/test_model.py                 | 9 +++++----
 4 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/flash/core/data/utils.py b/flash/core/data/utils.py
index 1402414d01..0155931398 100644
--- a/flash/core/data/utils.py
+++ b/flash/core/data/utils.py
@@ -156,7 +156,7 @@ def extract_tarfile(file_path: str, extract_path: str, mode: str):
                     try:
                         tar_ref.extract(member, path=extract_path, set_attrs=False)
                     except PermissionError:
-                        pass
+                        raise PermissionError("Could not extract tar file {}".format(file_path))
 
     if ".zip" in local_filename:
         if os.path.exists(local_filename):
diff --git a/flash/image/face_detection/backbones/fastface_backbones.py b/flash/image/face_detection/backbones/fastface_backbones.py
index 57062ab130..a8829d516b 100644
--- a/flash/image/face_detection/backbones/fastface_backbones.py
+++ b/flash/image/face_detection/backbones/fastface_backbones.py
@@ -24,7 +24,7 @@
     _MODEL_NAMES = []
 
 
-def fastface_backbone(model_name, pretrained, **kwargs):
+def fastface_backbone(model_name: str, pretrained: bool, **kwargs):
     if pretrained:
         pl_model = ff.FaceDetector.from_pretrained(model_name, **kwargs)
     else:
diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index 748a44292b..e2ceed03b4 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -79,7 +79,7 @@ def __init__(
         if model in ff.list_pretrained_models():
             model = FaceDetector.get_model(model, pretrained, **kwargs)
         else:
-            ValueError(f"{model} is not supported yet.")
+            ValueError(model + " is not supported yet, please select one from {}".format(ff.list_pretrained_models()))
 
         super().__init__(
             model=model,
@@ -93,8 +93,8 @@ def __init__(
 
     @staticmethod
     def get_model(
-        model_name,
-        pretrained,
+        model_name: str,
+        pretrained: bool,
         **kwargs,
     ):
         model, pl_model = FACE_DETECTION_BACKBONES.get(model_name)(pretrained=pretrained, **kwargs)
diff --git a/tests/image/face_detection/test_model.py b/tests/image/face_detection/test_model.py
index abdec7d8ae..e048cd353b 100644
--- a/tests/image/face_detection/test_model.py
+++ b/tests/image/face_detection/test_model.py
@@ -29,13 +29,14 @@
 
 
 @pytest.mark.skipif(not _FASTFACE_AVAILABLE, reason="fastface not installed.")
-def test_fastface_training():
+@pytest.mark.parametrize("model_name", ["lffd_slim", "lffd_original"])
+def test_fastface_training(tmpdir, model_name):
     dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="val")
-    datamodule = FaceDetectionData.from_datasets(train_dataset=dataset, batch_size=2)
+    datamodule = FaceDetectionData.from_datasets(train_dataset=dataset, batch_size=1)
 
-    model = FaceDetector(model="lffd_slim")
+    model = FaceDetector(model=model_name)
 
-    trainer = flash.Trainer(max_steps=2, num_sanity_val_steps=0)
+    trainer = flash.Trainer(max_steps=1, num_sanity_val_steps=0)
     trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
 

From 934d44be16ffb862cc074416ef3d512700d6cfb3 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 30 Sep 2021 13:35:50 +0000
Subject: [PATCH 32/35] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 flash/core/data/utils.py            | 2 +-
 flash/image/face_detection/model.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/flash/core/data/utils.py b/flash/core/data/utils.py
index 0155931398..c1bc18f698 100644
--- a/flash/core/data/utils.py
+++ b/flash/core/data/utils.py
@@ -156,7 +156,7 @@ def extract_tarfile(file_path: str, extract_path: str, mode: str):
                     try:
                         tar_ref.extract(member, path=extract_path, set_attrs=False)
                     except PermissionError:
-                        raise PermissionError("Could not extract tar file {}".format(file_path))
+                        raise PermissionError(f"Could not extract tar file {file_path}")
 
     if ".zip" in local_filename:
         if os.path.exists(local_filename):
diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index e2ceed03b4..5d7c6e0445 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -79,7 +79,7 @@ def __init__(
         if model in ff.list_pretrained_models():
             model = FaceDetector.get_model(model, pretrained, **kwargs)
         else:
-            ValueError(model + " is not supported yet, please select one from {}".format(ff.list_pretrained_models()))
+            ValueError(model + f" is not supported yet, please select one from {ff.list_pretrained_models()}")
 
         super().__init__(
             model=model,

From 164de7c00cdc98c70f68ceba17a197a3dbd0f77c Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Thu, 30 Sep 2021 09:36:01 -0400
Subject: [PATCH 33/35] .

---
 flash/core/data/utils.py            | 2 +-
 flash/image/face_detection/model.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/flash/core/data/utils.py b/flash/core/data/utils.py
index 0155931398..c1bc18f698 100644
--- a/flash/core/data/utils.py
+++ b/flash/core/data/utils.py
@@ -156,7 +156,7 @@ def extract_tarfile(file_path: str, extract_path: str, mode: str):
                     try:
                         tar_ref.extract(member, path=extract_path, set_attrs=False)
                     except PermissionError:
-                        raise PermissionError("Could not extract tar file {}".format(file_path))
+                        raise PermissionError(f"Could not extract tar file {file_path}")
 
     if ".zip" in local_filename:
         if os.path.exists(local_filename):
diff --git a/flash/image/face_detection/model.py b/flash/image/face_detection/model.py
index e2ceed03b4..5d7c6e0445 100644
--- a/flash/image/face_detection/model.py
+++ b/flash/image/face_detection/model.py
@@ -79,7 +79,7 @@ def __init__(
         if model in ff.list_pretrained_models():
             model = FaceDetector.get_model(model, pretrained, **kwargs)
         else:
-            ValueError(model + " is not supported yet, please select one from {}".format(ff.list_pretrained_models()))
+            ValueError(model + f" is not supported yet, please select one from {ff.list_pretrained_models()}")
 
         super().__init__(
             model=model,

From 538400370f720fcc82d897f112585a3ba5704a8c Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Thu, 30 Sep 2021 10:38:34 -0400
Subject: [PATCH 34/35] .

---
 tests/image/face_detection/test_model.py | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/tests/image/face_detection/test_model.py b/tests/image/face_detection/test_model.py
index e048cd353b..29f7206ac5 100644
--- a/tests/image/face_detection/test_model.py
+++ b/tests/image/face_detection/test_model.py
@@ -13,7 +13,9 @@
 # limitations under the License.
 import pytest
 
+import torch
 import flash
+
 from flash.core.registry import FlashRegistry
 from flash.core.utilities.imports import _FASTFACE_AVAILABLE
 from flash.image import FaceDetectionData, FaceDetector
@@ -29,22 +31,31 @@
 
 
 @pytest.mark.skipif(not _FASTFACE_AVAILABLE, reason="fastface not installed.")
-@pytest.mark.parametrize("model_name", ["lffd_slim", "lffd_original"])
-def test_fastface_training(tmpdir, model_name):
+def test_fastface_training():
     dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="val")
-    datamodule = FaceDetectionData.from_datasets(train_dataset=dataset, batch_size=1)
+    datamodule = FaceDetectionData.from_datasets(train_dataset=dataset, batch_size=2)
 
-    model = FaceDetector(model=model_name)
+    model = FaceDetector(model="lffd_slim")
 
-    trainer = flash.Trainer(max_steps=1, num_sanity_val_steps=0)
+    # test fit
+    trainer = flash.Trainer(max_steps=2, num_sanity_val_steps=0)
     trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
 
+@pytest.mark.skipif(not _FASTFACE_AVAILABLE, reason="fastface not installed.")
+def test_fastface_forward():
+    model = FaceDetector(model="lffd_slim")
+    mock_batch = torch.randn(2, 3, 256, 256)
+
+    # test model forward (tests: _prepare_batch, logits_to_preds, _postprocess from ff)
+    model(mock_batch)
+
+
 @pytest.mark.skipif(not _FASTFACE_AVAILABLE, reason="fastface not installed.")
 def test_fastface_backbones_registry():
     backbones = FACE_DETECTION_BACKBONES.available_keys()
     assert "lffd_slim" in backbones
     assert "lffd_original" in backbones
 
-    backbone, _ = FACE_DETECTION_BACKBONES.get("lffd_original")(pretrained=True)
+    backbone, _ = FACE_DETECTION_BACKBONES.get("lffd_original")(pretrained=False)
     assert isinstance(backbone, LFFD)

From a94b10737110b46ea247f91cfa90a2618a1a0899 Mon Sep 17 00:00:00 2001
From: ananyahjha93 <ananya@pytorchlightning.ai>
Date: Thu, 30 Sep 2021 10:41:21 -0400
Subject: [PATCH 35/35] .

---
 tests/image/face_detection/test_model.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/image/face_detection/test_model.py b/tests/image/face_detection/test_model.py
index 29f7206ac5..05228c6586 100644
--- a/tests/image/face_detection/test_model.py
+++ b/tests/image/face_detection/test_model.py
@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import pytest
-
 import torch
-import flash
 
+import flash
 from flash.core.registry import FlashRegistry
 from flash.core.utilities.imports import _FASTFACE_AVAILABLE
 from flash.image import FaceDetectionData, FaceDetector