Эх сурвалжийг харах

Separating weights into artistic, stable and video. Updating weight links throughout and adding them in readme

Jason Antic 6 жил өмнө
parent
commit
6563d38be7

+ 1 - 10
ImageColorizerArtistic.ipynb

@@ -804,7 +804,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "vis.plot_transformed_image(\"test_images/BritishSlum.jpg\", render_factor=47)"
+    "vis.plot_transformed_image(\"test_images/BritishSlum.jpg\", render_factor=45)"
    ]
   },
   {
@@ -1446,15 +1446,6 @@
     "vis.plot_transformed_image(\"test_images/1896NewsBoyGirl.jpg\", render_factor=21)"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "vis.plot_transformed_image(\"video/bwframes/video15/00320.jpg\", render_factor=28)"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,

+ 3 - 3
ImageColorizerColab.ipynb

@@ -185,7 +185,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "artistic = False #@param {type:\"boolean\"}"
+    "artistic = True #@param {type:\"boolean\"}"
    ]
   },
   {
@@ -196,9 +196,9 @@
    "source": [
     "!mkdir 'models'\n",
     "if artistic:\n",
-    "    !wget https://www.dropbox.com/s/9ne9su2mc5t0m38/ColorizeImagesArtistic_gen.pth?dl=0 -O ./models/ColorizeImagesArtistic_gen.pth\n",
+    "    !wget https://www.dropbox.com/s/zkehq1uwahhbc2o/ColorizeArtistic_gen.pth?dl=0 -O ./models/ColorizeArtistic_gen.pth\n",
     "else:\n",
-    "    !wget https://www.dropbox.com/s/ztgygpaz1z3jkjg/ColorizeImagesStable_gen.pth?dl=0 -O ./models/ColorizeImagesStable_gen.pth"
+    "    !wget https://www.dropbox.com/s/mwjep3vyqk5mkjc/ColorizeStable_gen.pth?dl=0 -O ./models/ColorizeStable_gen.pth"
    ]
   },
   {

+ 1 - 1
ImageColorizerStable.ipynb

@@ -237,7 +237,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "vis.plot_transformed_image(\"test_images/HPLabelleOfficeMontreal.jpg\", render_factor=45)"
+    "vis.plot_transformed_image(\"test_images/HPLabelleOfficeMontreal.jpg\", render_factor=44)"
    ]
   },
   {

+ 35 - 3
README.md

@@ -167,11 +167,13 @@ This project is built around the wonderful Fast.AI library.  Prereqs, in summary
 --------------------------
 #### Pretrained Weights 
 
-To start right away on your own machine with your own images or videos without training the models yourself, you'll need to download the weights and drop them in the /models/ folder.
+To start right away on your own machine with your own images or videos without training the models yourself, you'll need to download the generator weights and drop them in the /models/ folder.
 
-[Download image weights here](https://www.dropbox.com/s/3e4dqky91h824ik/ColorizeImages_gen.pth)
+[Download 'artistic' model weights here](https://www.dropbox.com/s/zkehq1uwahhbc2o/ColorizeArtistic_gen.pth?dl=0)
 
-[Download video weights here](https://www.dropbox.com/s/3e4dqky91h824ik/ColorizeVideos_gen.pth)
+[Download 'stable' model weights here](https://www.dropbox.com/s/mwjep3vyqk5mkjc/ColorizeStable_gen.pth?dl=0)
+
+[Download 'video' model weights here](https://www.dropbox.com/s/336vn9y4qwyg9yz/ColorizeVideo_gen.pth?dl=0)
 
 
 You can then do image colorization in this notebook:  [ImageColorizer.ipynb](ImageColorizer.ipynb) 
@@ -180,6 +182,8 @@ And you can do video colorization in this notebook:  [VideoColorizer.ipynb](Vide
 
 The notebooks should be able to guide you from here.
 
+A more complete list of available pretrained weights, including those for critics, can be found at the bottom of the readme.
+
 -------------------------
 ### Stuff That Should Probably Be In A Paper
 
@@ -210,7 +214,35 @@ Keep in mind- I haven't been entirely rigorous in figuring out what all is going
 
 This builds upon a technique developed in collaboration with Jeremy Howard and Sylvain Gugger for Fast.AI's Lesson 7 in version 3 of Practical Deep Learning for Coders Part I. The particular lesson notebook can be found here: https://github.com/fastai/course-v3/blob/master/nbs/dl1/lesson7-superres-gan.ipynb  
 
+-------------------------
+### Weights, Weights, and More Weights
+
+
+#### Completed Generator Weights
+
+[artistic](https://www.dropbox.com/s/zkehq1uwahhbc2o/ColorizeArtistic_gen.pth?dl=0)
+[stable](https://www.dropbox.com/s/mwjep3vyqk5mkjc/ColorizeStable_gen.pth?dl=0)
+[video](https://www.dropbox.com/s/336vn9y4qwyg9yz/ColorizeVideo_gen.pth?dl=0)
+
+#### Completed Critic Weights
 
+[artistic](https://www.dropbox.com/s/8g5txfzt2fw8mf5/ColorizeArtistic_crit.pth?dl=0)
+[stable](https://www.dropbox.com/s/7a8u20e7xdu1dtd/ColorizeStable_crit.pth?dl=0)
+[video](https://www.dropbox.com/s/0401djgo1dfxdzt/ColorizeVideo_crit.pth?dl=0)
+
+#### Pretrain Only Generator Weights
+
+[artistic](https://www.dropbox.com/s/9zexurvrve141n9/ColorizeArtistic_PretrainOnly_gen.pth?dl=0)
+[stable](https://www.dropbox.com/s/mdnuo1563bb8nh4/ColorizeStable_PretrainOnly_gen.pth?dl=0)
+[video](https://www.dropbox.com/s/avzixh1ujf86e8x/ColorizeVideo_PretrainOnly_gen.pth?dl=0)
+
+#### Pretrain Only Critic Weights
+
+[artistic](https://www.dropbox.com/s/lakxe8akzjgjnmh/ColorizeArtistic_PretrainOnly_crit.pth?dl=0)
+[stable](https://www.dropbox.com/s/b3wka56iyv1fvdc/ColorizeStable_PretrainOnly_crit.pth?dl=0)
+[video](https://www.dropbox.com/s/j7og84cbhpa94gs/ColorizeVideo_PretrainOnly_crit.pth?dl=0)
+
+-------------------------
 ### Want More?
 
 I'll be posting more results on Twitter. [<img src="result_images/Twitter_Social_Icon_Rounded_Square_Color.svg" width="28">](https://twitter.com/citnaj)

+ 1 - 1
VideoColorizerColab.ipynb

@@ -165,7 +165,7 @@
    "outputs": [],
    "source": [
     "!mkdir 'models'\n",
-    "!wget https://www.dropbox.com/s/ztgygpaz1z3jkjg/ColorizeImagesStable_gen.pth?dl=0 -O ./models/ColorizeImagesStable_gen.pth"
+    "!wget https://www.dropbox.com/s/336vn9y4qwyg9yz/ColorizeVideo_gen.pth?dl=0 -O ./models/ColorizeVideo_gen.pth"
    ]
   },
   {

+ 5 - 6
fasterai/visualize.py

@@ -21,7 +21,7 @@ class ModelImageVisualizer():
     
     def _clean_mem(self):
         return
-        #torch.cuda.empty_cache()
+        torch.cuda.empty_cache()
         #gc.collect()
 
     def _open_pil_image(self, path:Path)->Image:
@@ -145,28 +145,27 @@ class VideoColorizer():
 def get_video_colorizer(render_factor:int=36)->VideoColorizer:
     return get_stable_video_colorizer(render_factor=render_factor)
 
-def get_stable_video_colorizer(root_folder:Path=Path('./'), weights_name:str='ColorizeImagesStable_gen', 
+def get_stable_video_colorizer(root_folder:Path=Path('./'), weights_name:str='ColorizeVideo_gen', 
         results_dir='result_images', render_factor:int=36)->VideoColorizer:
     learn = gen_inference_wide(root_folder=root_folder, weights_name=weights_name)
     filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
     vis = ModelImageVisualizer(filtr, results_dir=results_dir)
     return VideoColorizer(vis)
 
-
-def get_image_colorizer(render_factor:int=36, artistic:bool=False)->ModelImageVisualizer:
+def get_image_colorizer(render_factor:int=36, artistic:bool=True)->ModelImageVisualizer:
     if artistic:
         return get_artistic_image_colorizer(render_factor=render_factor)
     else:
         return get_stable_image_colorizer(render_factor=render_factor)
 
-def get_stable_image_colorizer(root_folder:Path=Path('./'), weights_name:str='ColorizeImagesStable_gen', 
+def get_stable_image_colorizer(root_folder:Path=Path('./'), weights_name:str='ColorizeStable_gen', 
         results_dir='result_images', render_factor:int=36)->ModelImageVisualizer:
     learn = gen_inference_wide(root_folder=root_folder, weights_name=weights_name)
     filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
     vis = ModelImageVisualizer(filtr, results_dir=results_dir)
     return vis
 
-def get_artistic_image_colorizer(root_folder:Path=Path('./'), weights_name:str='ColorizeImagesArtistic_gen', 
+def get_artistic_image_colorizer(root_folder:Path=Path('./'), weights_name:str='ColorizeArtistic_gen', 
         results_dir='result_images', render_factor:int=36)->ModelImageVisualizer:
     learn = gen_inference_deep(root_folder=root_folder, weights_name=weights_name)
     filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)