Merge pull request #30 from luigi311/dev

Fix download models

Dockerfile

@@ -2,20 +2,18 @@ FROM docker.io/luigi311/low-power-image-processing-base-image:latest
 
 WORKDIR /app
 
-COPY download_models.sh entrypoint.sh ./
-
-# Fix return
-RUN sed -i 's/\r$//' download_models.sh && \
-    sed -i 's/\r$//' entrypoint.sh
-
-RUN chmod +x download_models.sh && ./download_models.sh
-
 COPY requirements.txt .
 
 # Install dependencies as root
 RUN pip install -r requirements.txt
 
 COPY . .
 
+# Fix return
+RUN sed -i 's/\r$//' download_models.sh && \
+    sed -i 's/\r$//' entrypoint.sh
+
+RUN chmod +x download_models.sh && ./download_models.sh
+
 # Entrypoint entrypoint.sh
 ENTRYPOINT ["/bin/bash", "entrypoint.sh"]

download_models.sh

@@ -8,8 +8,8 @@ declare -a image_adaptive_3dlut_models=( "LUTs" "LUTs_unpaired" "classifier" "cl
 mkdir -p super_resolution/ denoise/ffdnet denoise/ircnn color/image_adaptive_3dlut/pretrained_models/sRGB
 
 for scale in "${scales[@]}"; do
-    wget -O "super_resolution/ESPCN_x${scale}.pb" "https://raw.githubusercontent.com/fannymonori/TF-ESPCN/master/export/ESPCN_x${scale}.pb"
-    wget -O "super_resolution/FSRCNN_x${scale}.pb" "https://raw.githubusercontent.com/Saafke/FSRCNN_Tensorflow/master/models/FSRCNN_x${scale}.pb"
+    wget -O "super_resolution/opencv/ESPCN_x${scale}.pb" "https://raw.githubusercontent.com/fannymonori/TF-ESPCN/master/export/ESPCN_x${scale}.pb"
+    wget -O "super_resolution/opencv/FSRCNN_x${scale}.pb" "https://raw.githubusercontent.com/Saafke/FSRCNN_Tensorflow/master/models/FSRCNN_x${scale}.pb"
 done
 
 for model in "${ffdnet_models[@]}"; do

super_resolution/opencv/opencv_super_resolution.py

@@ -9,23 +9,23 @@ def opencv_super_resolution(image, method, scale):
     sr = cv2.dnn_superres.DnnSuperResImpl_create()
 
     if method == "ESPCN":
-        path = f"{model_path}/ESPCN_x{scale}.pb"
+        model_file = f"{model_path}/ESPCN_x{scale}.pb"
         url = f"https://raw.githubusercontent.com/fannymonori/TF-ESPCN/master/export/ESPCN_x{scale}.pb"
         model = "espcn"
     elif method == "FSRCNN":
-        path = f"{model_path}/FSRCNN_x{scale}.pb"
+        model_file = f"{model_path}/FSRCNN_x{scale}.pb"
         url = f"https://raw.githubusercontent.com/Saafke/FSRCNN_Tensorflow/master/models/FSRCNN_x{scale}.pb"
         model = "fsrcnn"
     else:
         raise Exception("Super Resolution: Method not supported")
 
     # If path does not exist, download the model
-    if not os.path.exists(path):
+    if not os.path.exists(model_file):
         print("Downloading model...")
-        downloader(url, path, model_path)
+        downloader(url, model_file, model_path)
         print("Download Complete")
 
-    sr.readModel(path)
+    sr.readModel(model_file)
     sr.setModel(model, scale)
     print("Running Super Sampling")
     result = sr.upsample(image)