Update README.md

README.md

@@ -8,7 +8,9 @@ pipeline_tag: image-classification
 ---
 # Google Safesearch Mini Model Card
 
-Initially, the training data consisted of 278,000 images, and the model achieved 99% training and test acc. Now, this model is trained on 2,220,000+ images scraped from Google Images, Reddit, Imgur, and Github.
+<a href="https://huggingface.co/FredZhang7/google-safesearch-mini-v2"> <font size="4"> <bold> Version 2 is here! </bold> </font> </a>
+
+This model is trained on 2,220,000+ images scraped from Google Images, Reddit, Imgur, and Github.
 The InceptionV3 and Xception models have been fine-tuned to predict the likelihood of an image falling into one of three categories: nsfw_gore, nsfw_suggestive, and safe.
 
 After 20 epochs on PyTorch, the finetuned InceptionV3 model achieves 94% acc on both training and test data. After 3.3 epochs on Keras, the finetuned Xception model scores 94% acc on training set and 92% on test set.
@@ -23,64 +25,46 @@ The PyTorch model runs much slower with transformers, so downloading it external
 pip install --upgrade torchvision
 ```
 ```python
-import torch, os
+import torch, os, warnings, requests
+from io import BytesIO
 from PIL import Image
-import warnings
-warnings.filterwarnings("ignore")
+from urllib.request import urlretrieve
+from torchvision import transforms
 
 PATH_TO_IMAGE = 'https://images.unsplash.com/photo-1594568284297-7c64464062b1'
 USE_CUDA = False
 
+warnings.filterwarnings("ignore")
 def download_model():
     print("Downloading google_safesearch_mini.bin...")
-    import urllib.request
-    url = "https://huggingface.co/FredZhang7/google-safesearch-mini/resolve/main/pytorch_model.bin"
-    urllib.request.urlretrieve(url, "google_safesearch_mini.bin")
+    urlretrieve("https://huggingface.co/FredZhang7/google-safesearch-mini/resolve/main/pytorch_model.bin", "google_safesearch_mini.bin")
 
-def run():
+def eval():
     if not os.path.exists("google_safesearch_mini.bin"):
         download_model()
     model = torch.jit.load('./google_safesearch_mini.bin')
-    if PATH_TO_IMAGE.startswith('http://') or PATH_TO_IMAGE.startswith('https://'):
-        import requests
-        from io import BytesIO
-        response = requests.get(PATH_TO_IMAGE)
-        img = Image.open(BytesIO(response.content)).convert('RGB')
-    else:
-        img = Image.open(PATH_TO_IMAGE).convert('RGB')
-    from torchvision import transforms
-    transform = transforms.Compose([
-        transforms.Resize(299),
-        transforms.ToTensor(),
-        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
-    ])
-    img = transform(img)
-    img = img.unsqueeze(0)
+    img = Image.open(PATH_TO_IMAGE).convert('RGB') if not (PATH_TO_IMAGE.startswith('http://') or PATH_TO_IMAGE.startswith('https://')) else Image.open(BytesIO(requests.get(PATH_TO_IMAGE).content)).convert('RGB')
+    transform = transforms.Compose([transforms.Resize(299), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+    img = transform(img).unsqueeze(0)
     if USE_CUDA:
-        img = img.cuda()
-        model = model.cuda()
+        img, model = img.cuda(), model.cuda()
     else:
-        img = img.cpu()
-        model = model.cpu()
+        img, model = img.cpu(), model.cpu()
     model.eval()
     with torch.no_grad():
         out, _ = model(img)
         _, predicted = torch.max(out.data, 1)
-        classes = {
-            0: 'nsfw_gore',
-            1: 'nsfw_suggestive',
-            2: 'safe'
-        }
+        classes = {0: 'nsfw_gore', 1: 'nsfw_suggestive', 2: 'safe'}
+
         # account for edge cases
-        if predicted[0] != 2 and abs(out[0][2] - out[0][predicted[0]]) > 0.22:
+        if predicted[0] != 2 and abs(out[0][2] - out[0][predicted[0]]) > 0.20:
             img = Image.new('RGB', image.size, color = (0, 255, 255))
             print("\033[93m" + "safe" + "\033[0m")
         else:
             print('\n\033[1;31m' + classes[predicted.item()] + '\033[0m' if predicted.item() != 2 else '\033[1;32m' + classes[predicted.item()] + '\033[0m\n')
 
 if __name__ == '__main__':
-    torch.multiprocessing.freeze_support()
-    run()
+    eval()
 ```
 Output Example:
 ![prediction](./output_example.png)