From 0189f009d9d0f3e6a2aa4f83e15223847f67fd00 Mon Sep 17 00:00:00 2001
From: kshitijrajsharma <skshitizraj@gmail.com>
Date: Tue, 29 Aug 2023 14:01:59 +0545
Subject: [PATCH] Raise error if there are no images in the area

---
 backend/core/tasks.py | 10 +++++++++-
 backend/core/utils.py | 14 +++++++++++++-
 backend/core/views.py |  3 +++
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/backend/core/tasks.py b/backend/core/tasks.py
index 4cd34642..0ebf7c1d 100644
--- a/backend/core/tasks.py
+++ b/backend/core/tasks.py
@@ -16,7 +16,12 @@
     FeedbackLabelFileSerializer,
     LabelFileSerializer,
 )
-from core.utils import bbox, download_imagery, get_start_end_download_coords
+from core.utils import (
+    bbox,
+    download_imagery,
+    get_start_end_download_coords,
+    is_dir_empty,
+)
 from django.conf import settings
 from django.contrib.gis.db.models.aggregates import Extent
 from django.contrib.gis.geos import GEOSGeometry
@@ -106,8 +111,11 @@ def train_model(
                             base_path=training_input_image_source,
                             source=source_imagery,
                         )
+
                     except Exception as ex:
                         raise ex
+                if is_dir_empty(training_input_image_source):
+                    raise ValueError("No images found in the area")
 
             ## -----------LABEL GENERATOR---------
             logging.info("Label Generator started")
diff --git a/backend/core/utils.py b/backend/core/utils.py
index fdc3fe32..b23dd3b5 100644
--- a/backend/core/utils.py
+++ b/backend/core/utils.py
@@ -115,12 +115,24 @@ def get_start_end_download_coords(bbox_coords, zm_level, tile_size):
 import logging
 
 
+def is_dir_empty(directory_path):
+    return not any(os.scandir(directory_path))
+
+
 def download_image(url, base_path, source_name):
     response = requests.get(url)
+
+    # Check if the URL request was successful
+    if response.status_code != 200:
+        raise ValueError(
+            f"Failed to download image from {url}. Status code: {response.status_code}"
+        )
+
     image = response.content
-    url = re.sub(r"\.(png|jpeg)$", "", url)
 
+    url = re.sub(r"\.(png|jpeg)$", "", url)
     url_splitted_list = url.split("/")
+
     filename = f"{base_path}/{source_name}-{url_splitted_list[-2]}-{url_splitted_list[-1]}-{url_splitted_list[-3]}.png"
 
     with open(filename, "wb") as f:
diff --git a/backend/core/views.py b/backend/core/views.py
index cb20ab9a..ecb25da0 100644
--- a/backend/core/views.py
+++ b/backend/core/views.py
@@ -68,6 +68,7 @@
     get_dir_size,
     get_start_end_download_coords,
     gpx_generator,
+    is_dir_empty,
     process_rawdata,
     request_rawdata,
 )
@@ -532,6 +533,8 @@ def post(self, request, *args, **kwargs):
                 )
                 prediction_output = f"{temp_path}/prediction/output"
                 print("Image Downloaded , Starting Inference")
+                if is_dir_empty(temp_path):
+                    return Response("No Images found", status=500)
                 start_time = time.time()
                 model_path = os.path.join(
                     settings.TRAINING_WORKSPACE,