Refactor load_image and download_file to check if file is already exists

openvinotoolkit · Nov 28, 2024 · 86bfa5f · 86bfa5f
1 parent d4ac211
commit 86bfa5f
Show file tree

Hide file tree

Showing 10 changed files with 63 additions and 92 deletions.
diff --git a/notebooks/depth-anything/depth-anything-v2.ipynb b/notebooks/depth-anything/depth-anything-v2.ipynb
@@ -197,11 +197,10 @@
     "from notebook_utils import download_file, device_widget, quantization_widget\n",
     "\n",
     "\n",
-    "if not Path(\"furseal.png\").exists():\n",
-    "    download_file(\n",
-    "        \"https://github.com/openvinotoolkit/openvino_notebooks/assets/29454499/3f779fc1-c1b2-4dec-915a-64dae510a2bb\",\n",
-    "        \"furseal.png\",\n",
-    "    )\n",
+    "download_file(\n",
+    "    \"https://github.com/openvinotoolkit/openvino_notebooks/assets/29454499/3f779fc1-c1b2-4dec-915a-64dae510a2bb\",\n",
+    "    \"furseal.png\",\n",
+    ")\n",
     "\n",
     "Image.open(\"furseal.png\").resize((600, 400))"
    ]
@@ -547,11 +546,10 @@
    "source": [
     "VIDEO_FILE = \"./Coco Walking in Berkeley.mp4\"\n",
     "\n",
-    "if not Path(VIDEO_FILE).exists():\n",
-    "    download_file(\n",
-    "        \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/Coco%20Walking%20in%20Berkeley.mp4\",\n",
-    "        VIDEO_FILE,\n",
-    "    )\n",
+    "download_file(\n",
+    "    \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/Coco%20Walking%20in%20Berkeley.mp4\",\n",
+    "    VIDEO_FILE,\n",
+    ")\n",
     "\n",
     "# Number of seconds of input video to process. Set `NUM_SECONDS` to 0 to process\n",
     "# the full video.\n",

diff --git a/notebooks/optical-character-recognition/optical-character-recognition.ipynb b/notebooks/optical-character-recognition/optical-character-recognition.ipynb
@@ -261,7 +261,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "id": "66495015",
    "metadata": {
     "tags": []
@@ -282,7 +282,7 @@
     "# The `image_file` variable can point to a URL or a local image.\n",
     "image_file = \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/intel_rnb.jpg\"\n",
     "\n",
-    "image = load_image(image_file)\n",
+    "image = load_image(\"intel_rnb.jpg\", image_file)\n",
     "\n",
     "# N,C,H,W = batch size, number of channels, height, width.\n",
     "N, C, H, W = detection_input_layer.shape\n",

diff --git a/notebooks/tflite-selfie-segmentation/tflite-selfie-segmentation.ipynb b/notebooks/tflite-selfie-segmentation/tflite-selfie-segmentation.ipynb
@@ -146,8 +146,7 @@
     "tflite_model_path = Path(\"selfie_multiclass_256x256.tflite\")\n",
     "tflite_model_url = \"https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_multiclass_256x256/float32/latest/selfie_multiclass_256x256.tflite\"\n",
     "\n",
-    "if not tflite_model_path.exists():\n",
-    "    download_file(tflite_model_url, tflite_model_path)"
+    "download_file(tflite_model_url, tflite_model_path)"
    ]
   },
   {
@@ -337,12 +336,9 @@
     "\n",
     "# Read input image and convert it to RGB\n",
     "test_image_url = \"https://user-images.githubusercontent.com/29454499/251036317-551a2399-303e-4a4a-a7d6-d7ce973e05c5.png\"\n",
-    "image_path = \"example-img.png\"\n",
+    "image_name = \"example-img.png\"\n",
     "\n",
-    "if not Path(image_path).exists():\n",
-    "    img = download_file(test_image_url, image_path)\n",
-    "\n",
-    "img = load_image(image_path)\n",
+    "img = load_image(image_name, test_image_url)\n",
     "img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
     "\n",
     "\n",
@@ -712,11 +708,10 @@
     "    VIDEO_SOURCE = 0  # Webcam\n",
     "else:\n",
     "    VIDEO_SOURCE = \"CEO-Pat-Gelsinger-on-Leading-Intel.mp4\"\n",
-    "    if not Path(VIDEO_SOURCE).exists():\n",
-    "        download_file(\n",
-    "            \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/CEO%20Pat%20Gelsinger%20on%20Leading%20Intel.mp4\",\n",
-    "            VIDEO_SOURCE,\n",
-    "        )"
+    "    download_file(\n",
+    "        \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/CEO%20Pat%20Gelsinger%20on%20Leading%20Intel.mp4\",\n",
+    "        VIDEO_SOURCE,\n",
+    "    )"
    ]
   },
   {

diff --git a/notebooks/tflite-to-openvino/tflite-to-openvino.ipynb b/notebooks/tflite-to-openvino/tflite-to-openvino.ipynb
@@ -196,14 +196,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "id": "07790e19-baec-4368-8650-c0381e2352e5",
    "metadata": {
     "tags": []
    },
    "outputs": [],
    "source": [
-    "image = load_image(\"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_bricks.png\")\n",
+    "image = load_image(\"coco_bricks.png\", \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_bricks.png\")\n",
     "# load_image reads the image in BGR format, [:,:,::-1] reshape transfroms it to RGB\n",
     "image = Image.fromarray(image[:, :, ::-1])\n",
     "resized_image = image.resize((224, 224))\n",

diff --git a/notebooks/tiny-sd-image-generation/tiny-sd-image-generation.ipynb b/notebooks/tiny-sd-image-generation/tiny-sd-image-generation.ipynb
@@ -1141,7 +1141,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1189,8 +1189,9 @@
     "from diffusers.utils import load_image\n",
     "\n",
     "default_image_url = \"https://user-images.githubusercontent.com/29454499/260418860-69cc443a-9ee6-493c-a393-3a97af080be7.jpg\"\n",
+    "image_name = \"default.jpg\"\n",
     "# read uploaded image\n",
-    "image = load_image(default_image_url)\n",
+    "image = load_image(image_name, default_image_url)\n",
     "print(\"Pipeline settings\")\n",
     "print(f\"Input positive prompt: \\n\\t{text_prompt_i2i}\")\n",
     "print(f\"Input negative prompt: \\n\\t{negative_prompt_i2i}\")\n",

diff --git a/notebooks/vision-background-removal/vision-background-removal.ipynb b/notebooks/vision-background-removal/vision-background-removal.ipynb
@@ -366,7 +366,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "id": "1a86e7fd",
    "metadata": {
     "colab": {
@@ -381,12 +381,13 @@
    "outputs": [],
    "source": [
     "IMAGE_URI = \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_hollywood.jpg\"\n",
+    "IMAGE_NAME = \"coco_hollywood.jpg\"\n",
     "\n",
     "input_mean = np.array([123.675, 116.28, 103.53]).reshape(1, 3, 1, 1)\n",
     "input_scale = np.array([58.395, 57.12, 57.375]).reshape(1, 3, 1, 1)\n",
     "\n",
     "image = cv2.cvtColor(\n",
-    "    src=load_image(IMAGE_URI),\n",
+    "    src=load_image(IMAGE_NAME, IMAGE_URI),\n",
     "    code=cv2.COLOR_BGR2RGB,\n",
     ")\n",
     "\n",
@@ -571,7 +572,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "id": "f5628a0c",
    "metadata": {
     "tags": [
@@ -603,12 +604,13 @@
     }
    ],
    "source": [
-    "BACKGROUND_FILE = \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/wall.jpg\"\n",
+    "BACKGROUND_IMAGE_URL = \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/wall.jpg\"\n",
+    "BACKGROUND_IMAGE_NAME = \"wall.jpg\"\n",
     "OUTPUT_DIR = \"output\"\n",
     "\n",
     "os.makedirs(name=OUTPUT_DIR, exist_ok=True)\n",
     "\n",
-    "background_image = cv2.cvtColor(src=load_image(BACKGROUND_FILE), code=cv2.COLOR_BGR2RGB)\n",
+    "background_image = cv2.cvtColor(src=load_image(BACKGROUND_IMAGE_NAME, BACKGROUND_IMAGE_URL), code=cv2.COLOR_BGR2RGB)\n",
     "background_image = cv2.resize(src=background_image, dsize=(image.shape[1], image.shape[0]))\n",
     "\n",
     "# Set all the foreground pixels from the result to 0\n",
@@ -617,7 +619,7 @@
     "new_image = background_image + bg_removed_result\n",
     "\n",
     "# Save the generated image.\n",
-    "new_image_path = Path(f\"{OUTPUT_DIR}/{Path(IMAGE_URI).stem}-{Path(BACKGROUND_FILE).stem}.jpg\")\n",
+    "new_image_path = Path(f\"{OUTPUT_DIR}/{Path(IMAGE_URI).stem}-{BACKGROUND_IMAGE_NAME}\")\n",
     "cv2.imwrite(filename=str(new_image_path), img=cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR))\n",
     "\n",
     "# Display the original image and the image with the new background side by side\n",

diff --git a/notebooks/vision-monodepth/vision-monodepth.ipynb b/notebooks/vision-monodepth/vision-monodepth.ipynb
@@ -172,11 +172,8 @@
     "ir_model_name_xml = \"MiDaS_small.xml\"\n",
     "ir_model_name_bin = \"MiDaS_small.bin\"\n",
     "\n",
-    "\n",
-    "if not (model_folder / ir_model_name_xml).exists():\n",
-    "    download_file(ir_model_url + ir_model_name_xml, filename=ir_model_name_xml, directory=model_folder)\n",
-    "if not (model_folder / ir_model_name_bin).exists():\n",
-    "    download_file(ir_model_url + ir_model_name_bin, filename=ir_model_name_bin, directory=model_folder)\n",
+    "download_file(ir_model_url + ir_model_name_xml, filename=ir_model_name_xml, directory=model_folder)\n",
+    "download_file(ir_model_url + ir_model_name_bin, filename=ir_model_name_bin, directory=model_folder)\n",
     "\n",
     "model_xml_path = model_folder / ir_model_name_xml"
    ]
@@ -348,12 +345,10 @@
    "outputs": [],
    "source": [
     "IMAGE_URL = \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_bike.jpg\"\n",
-    "IMAGE_PATH = \"coco_bike.jpg\"\n",
+    "IMAGE_NAME = \"coco_bike.jpg\"\n",
     "\n",
-    "if not Path(IMAGE_PATH).exists():\n",
-    "    download_file(IMAGE_URL, IMAGE_PATH)\n",
     "\n",
-    "image = load_image(IMAGE_PATH)\n",
+    "image = load_image(IMAGE_NAME, IMAGE_URL)\n",
     "\n",
     "\n",
     "# Resize to input shape for network.\n",
@@ -476,11 +471,10 @@
    "source": [
     "# Video source: https://www.youtube.com/watch?v=fu1xcQdJRws (Public Domain)\n",
     "VIDEO_FILE = \"Coco-Walking-in-Berkeley.mp4\"\n",
-    "if not Path(VIDEO_FILE).exists():\n",
-    "    download_file(\n",
-    "        \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/Coco%20Walking%20in%20Berkeley.mp4\",\n",
-    "        VIDEO_FILE,\n",
-    "    )\n",
+    "download_file(\n",
+    "    \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/Coco%20Walking%20in%20Berkeley.mp4\",\n",
+    "    VIDEO_FILE,\n",
+    ")\n",
     "\n",
     "# Number of seconds of input video to process. Set `NUM_SECONDS` to 0 to process\n",
     "# the full video.\n",

diff --git a/notebooks/yolov11-optimization/yolov11-instance-segmentation.ipynb b/notebooks/yolov11-optimization/yolov11-instance-segmentation.ipynb
@@ -155,12 +155,11 @@
     "# Download a test sample\n",
     "IMAGE_PATH = Path(\"./data/coco_bike.jpg\")\n",
     "\n",
-    "if not IMAGE_PATH.exists():\n",
-    "    download_file(\n",
-    "        url=\"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_bike.jpg\",\n",
-    "        filename=IMAGE_PATH.name,\n",
-    "        directory=IMAGE_PATH.parent,\n",
-    "    )"
+    "download_file(\n",
+    "    url=\"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_bike.jpg\",\n",
+    "    filename=IMAGE_PATH.name,\n",
+    "    directory=IMAGE_PATH.parent,\n",
+    ")"
    ]
   },
   {
@@ -1195,11 +1194,10 @@
     "if WEBCAM_INFERENCE:\n",
     "    VIDEO_SOURCE = 0  # Webcam\n",
     "else:\n",
-    "    if not Path(\"people.mp4\").exists():\n",
-    "        download_file(\n",
-    "            \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/people.mp4\",\n",
-    "            \"people.mp4\",\n",
-    "        )\n",
+    "    download_file(\n",
+    "        \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/people.mp4\",\n",
+    "        \"people.mp4\",\n",
+    "    )\n",
     "    VIDEO_SOURCE = \"people.mp4\""
    ]
   },

diff --git a/notebooks/yolov11-optimization/yolov11-object-detection.ipynb b/notebooks/yolov11-optimization/yolov11-object-detection.ipynb
@@ -156,12 +156,11 @@
     "# Download a test sample\n",
     "IMAGE_PATH = Path(\"./data/coco_bike.jpg\")\n",
     "\n",
-    "if not IMAGE_PATH.exists():\n",
-    "    download_file(\n",
-    "        url=\"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_bike.jpg\",\n",
-    "        filename=IMAGE_PATH.name,\n",
-    "        directory=IMAGE_PATH.parent,\n",
-    "    )"
+    "download_file(\n",
+    "    url=\"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_bike.jpg\",\n",
+    "    filename=IMAGE_PATH.name,\n",
+    "    directory=IMAGE_PATH.parent,\n",
+    ")"
    ]
   },
   {
@@ -1139,11 +1138,10 @@
     "if WEBCAM_INFERENCE:\n",
     "    VIDEO_SOURCE = 0  # Webcam\n",
     "else:\n",
-    "    if not Path(\"people.mp4\").exists():\n",
-    "        download_file(\n",
-    "            \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/people.mp4\",\n",
-    "            \"people.mp4\",\n",
-    "        )\n",
+    "    download_file(\n",
+    "        \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/people.mp4\",\n",
+    "        \"people.mp4\",\n",
+    "    )\n",
     "    VIDEO_SOURCE = \"people.mp4\""
    ]
   },

diff --git a/utils/notebook_utils.py b/utils/notebook_utils.py
@@ -79,41 +79,26 @@ def pip_install(*args):
     subprocess.run([sys.executable, "-m", "pip", "install", *cli_args], shell=(platform.system() == "Windows"), check=True)
 
 
-def load_image(path: str) -> np.ndarray:
+def load_image(name: str, url: str = None) -> np.ndarray:
     """
-    Loads an image from `path` and returns it as BGR numpy array. `path`
-    should point to an image file, either a local filename or a url. The image is
-    not stored to the filesystem. Use the `download_file` function to download and
-    store an image.
+    Loads an image by `url` and returns it as BGR numpy array. The image is
+    stored to the filesystem with name `name`. If the image file already exists
+    loads the local image.
 
-    :param path: Local path name or URL to image.
+    :param name: Local path name of the image.
+    :param url: url to the image
     :return: image as BGR numpy array
     """
     import cv2
     import requests
 
-    if path.startswith("http"):
-        # Set User-Agent to Mozilla because some websites block
-        # requests with User-Agent Python
-        response = requests.get(path, headers={"User-Agent": "Mozilla/5.0"})
-        array = np.asarray(bytearray(response.content), dtype="uint8")
-        image = cv2.imdecode(array, -1)  # Loads the image as BGR
-    else:
-        image = cv2.imread(path)
-    return image
-
-
-def download_image(name: str, url: str = None):
-    import cv2
-    import requests
-
     if not Path(name).exists():
         # Set User-Agent to Mozilla because some websites block
         # requests with User-Agent Python
         response = requests.get(url, headers={"User-Agent": "Mozilla/5.0"})
         array = np.asarray(bytearray(response.content), dtype="uint8")
         image = cv2.imdecode(array, -1)  # Loads the image as BGR
-        image.save(name)
+        cv2.imwrite(name, image)
     else:
         image = cv2.imread(name)