output in png

lqhl · Oct 5, 2023 · edcf6bc · edcf6bc
1 parent 5ca1078
commit edcf6bc
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/cog_predict.py b/cog_predict.py
@@ -123,11 +123,7 @@ def predict(
                 print('Error', error)
                 print('If you encounter CUDA out of memory, try to set "tile" to a smaller size, e.g., 400.')
 
-            if img_mode == 'RGBA':  # RGBA images should be saved in png format
-                extension = 'png'
-            # save_path = f'output/out.{extension}'
-            # cv2.imwrite(save_path, output)
-            out_path = Path(tempfile.mkdtemp()) / f'out.{extension}'
+            out_path = Path(tempfile.mkdtemp()) / f'out.png'
             cv2.imwrite(str(out_path), output)
         except Exception as error:
             print('global exception: ', error)