Adding rest of files

main.py

@@ -0,0 +1,76 @@
+import plantvision
+import requests
+from io import BytesIO
+import pickle as pkl
+from flask import Flask, render_template, request, session, jsonify, url_for
+from PIL import Image
+import os
+import time
+import random
+from pathlib import Path
+THIS_FOLDER = Path(__file__).parent.resolve()
+
+app = Flask(__name__)
+app.secret_key = 'pi-33pp-co-sk-33'
+app.template_folder = os.path.abspath(f'{THIS_FOLDER}/web/templates')
+app.static_folder = os.path.abspath(f'{THIS_FOLDER}/web/static')
+print(app.static_folder)
+
+flowerLayers = None
+leafLayers = None
+fruitLayers = None
+
+@app.route('/')
+def home():
+    return render_template('index.html')
+
+@app.route('/guess', methods=['POST'])
+def guess():
+    global flowerLayers, leafLayers, fruitLayers
+
+    if request.method == 'POST':
+        print('Thinking...')
+        
+        img = request.files.get('uploaded-image')
+        feature = request.form.get('feature')
+        
+        tensor = plantvision.processImage(img, feature)
+        predictions = plantvision.see(tensor, feature, 6)
+
+        with open(f'{THIS_FOLDER}/resources/speciesNameToKey.pkl','rb') as f:
+            speciesNameToKey = pkl.load(f)
+        with open(f'{THIS_FOLDER}/resources/speciesNameToVernacular.pkl','rb') as f:
+            speciesNameToVernacular = pkl.load(f)
+        with open(f'{THIS_FOLDER}/resources/{feature}speciesIndexDict.pkl','rb') as f:
+            speciesNameToIndex = pkl.load(f)
+
+        urls = []
+        predicted_image_urls = []
+        for p in predictions:
+            key = speciesNameToKey[p]
+            img = speciesNameToIndex[p]
+            query = ''
+            for i in p.split(' '):
+                query += i 
+                query += '+'
+            urls.append(f'https://www.google.com/search?q={query[:-1]}')
+            predicted_image_urls.append(f"https://storage.googleapis.com/bmllc-images-bucket/images/img{img}.jpeg")
+
+        names = []
+        for p in predictions:
+            try:
+                names.append(speciesNameToVernacular[p])
+            except:
+                names.append(p)
+
+        response = {
+            'names': names,
+            'species': predictions,
+            'predictions': urls,
+            'images': predicted_image_urls
+        }
+
+        return jsonify(response)
+
+if __name__ == '__main__':
+    app.run(port=int(os.environ.get("PORT", 8080)),host='0.0.0.0',debug=True)

plantvision.py

@@ -0,0 +1,115 @@
+import requests
+from io import BytesIO
+from PIL import Image, ImageOps
+import torchvision.transforms as T
+import torch
+import gc
+import pickle as pkl
+from pathlib import Path
+THIS_FOLDER = Path(__file__).parent.resolve()
+import datetime as dt
+from transformers import AutoModel
+import torch.nn as nn
+import torch.nn.functional as F
+import logging
+logging.disable(logging.INFO)
+logging.disable(logging.WARNING)
+
+visionTransformer = AutoModel.from_pretrained(r"google/vit-base-patch16-224-in21k")
+
+class PlantVision(nn.Module):
+    def __init__(self, num_classes):
+        super(PlantVision, self).__init__()
+        self.vit = visionTransformer
+        count = 0
+        for child in self.vit.children():
+            count += 1
+            if count < 4:
+                for param in child.parameters():
+                    param.requires_grad = False
+        self.vitLayers = list(self.vit.children()) 
+        self.vitTop = nn.Sequential(*self.vitLayers[:-2])
+        self.vitNorm = list(self.vit.children())[2]
+        self.vit = None
+        gc.collect()
+        self.vitFlatten = nn.Flatten()
+        self.vitLinear = nn.Linear(151296,num_classes)
+        self.fc = nn.Linear(num_classes, num_classes)
+
+    def forward(self, input):
+        output = self.vitTop(input).last_hidden_state
+        output = self.vitNorm(output)
+        output = self.vitFlatten(output)
+        output = F.relu(self.vitLinear(output))
+        output = self.fc(output)
+        return output
+
+device = 'cpu' # ('cuda' if torch.cuda.is_available else 'cpu')
+
+with open(fr'{THIS_FOLDER}/resources/flowerLabelSet.pkl', 'rb') as f:
+        flowerLabelSet = pkl.load(f)
+
+with open(fr'{THIS_FOLDER}/resources/leafLabelSet.pkl', 'rb') as f:
+        leafLabelSet = pkl.load(f)
+
+with open(fr'{THIS_FOLDER}/resources/fruitLabelSet.pkl', 'rb') as f:
+        fruitLabelSet = pkl.load(f)
+
+def loadModel(feature, labelSet):
+    model = PlantVision(num_classes=len(labelSet))
+    model.vitFlatten.load_state_dict(torch.load(BytesIO(requests.get(f"https://storage.googleapis.com/bmllc-plant-model-bucket/{feature}-vitFlatten-weights.pt").content), map_location=torch.device(device)), strict=False)
+    model.vitLinear.load_state_dict(torch.load(BytesIO(requests.get(f"https://storage.googleapis.com/bmllc-plant-model-bucket/{feature}-vitLinear-weights.pt").content), map_location=torch.device(device)), strict=False)
+    model.fc.load_state_dict(torch.load(BytesIO(requests.get(f"https://storage.googleapis.com/bmllc-plant-model-bucket/{feature}-fc-weights.pt").content), map_location=torch.device(device)), strict=False)
+    model = model.half()
+    return model
+
+start = dt.datetime.now()
+flower = loadModel('flower',flowerLabelSet)
+leaf = loadModel('leaf',leafLabelSet)
+fruit = loadModel('fruit',fruitLabelSet)
+print(dt.datetime.now() - start)
+
+def processImage(imagePath, feature):
+    with open(fr'{THIS_FOLDER}/resources/{feature}MeansAndStds.pkl', 'rb') as f:
+        meansAndStds = pkl.load(f)
+
+    img = Image.open(imagePath).convert('RGB')
+    cropped = ImageOps.fit(img, (224,224), Image.Resampling.LANCZOS)
+
+    process = T.Compose([
+        T.CenterCrop(224),
+        T.ToTensor(),
+        T.ConvertImageDtype(torch.float32),
+        T.Normalize(
+       mean=meansAndStds['mean'],
+       std=meansAndStds['std'])
+    ])
+
+    return process(cropped)
+
+def see(tensor,feature,k):
+
+        if feature=='flower':
+                model = flower.float()
+                labelSet = flowerLabelSet
+        
+        elif feature=='leaf':
+                model = leaf.float()
+                labelSet = leafLabelSet
+        
+        elif feature=='fruit':
+                model = fruit.float()
+                labelSet = fruitLabelSet
+        
+        with torch.no_grad():        
+                output = model(tensor.unsqueeze(0))
+                top = torch.topk(output,k,dim=1)
+                predictions = top.indices[0]
+        
+        predictedSpecies = []
+        for i in predictions:
+                predictedSpecies.append(labelSet[i])
+        
+        model = None
+        gc.collect()
+        return predictedSpecies

requirements.txt

@@ -0,0 +1,11 @@
+Flask==2.1.0
+gunicorn==20.1.0
+torch
+torchvision
+pillow
+tqdm
+argparse
+pathlib
+transformers
+fastapi
+uvicorn

resources/flowerImageIndices.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9189a92359774fb40bb605424ae3b3a2e6ae2b9a92b1c0bb61fb79981549ddfb
+size 113164

resources/flowerImageLabels.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9fd985d5e407b226c439e966aadc18f6332195f8cd118ccdd18bd4cf1de54b63
+size 203485

resources/flowerLabelSet.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:89ba3712eb686badc02a768e6b0f09974c4161466ea92cee1c9b86f70e33af15
+size 23964

resources/flowerMeansAndStds.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:76a3349b54abf757e5ef18d4801320c8194561eedab4ce3ea9739ff2bca25d8f
+size 236

resources/flowerspeciesIndexDict.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:571c2b76bc6f92e97bd51bd37224b62f7c1ecc6c02dd7768c2c1cc0212f9d26b
+size 70059

resources/fruitImageIndices.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5a902edb6214f6c4eedba3213b7bdd24c5f87cc62db33c66f645a4bd0444f0f
+size 17734

resources/fruitImageLabels.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c260c9bce0614a1197a92ffa29600d3890e97f733e702c4a5870a21eb8f665a8
+size 27866

resources/fruitLabelSet.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:875210d10918a499c84815b3393afc8f3e4fc07979c872714053c40059857da5
+size 5164

resources/fruitMeansAndStds.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1479a0b5e531db1b783fa3fe8f47bb6278fc43c0dc9448decb53f855df83fe72
+size 236

resources/fruitspeciesIndexDict.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b045988de9a625ce8ebb81adf64f00bf2709ce76868c388c95aeaeb1a918c930
+size 22526

resources/fruitspeciesNameToID.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ffb83cf9e6053dbc3888d5e6b3e2d784bfd09333fb5058b51f3f1da0d005fc1
+size 584518

resources/imageIndices.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:99ebbdcb3b5124aab5cb156b8be0672041fa16b70acb5e1efbb0a796dd54cf5b
+size 650154

resources/imageLabels.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:76c77a8a0bf3e1316e382b13db60c9b6dc90143ba9712f07f24ea6684b5b7da9
+size 1771457

resources/imageLocations.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3af966e13a46d4fddda1abe0764db5a0fa3b1996b113756b7fe2db1160a4ebbb
+size 2234643

resources/imageSelections.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c306952141e794e935a1bdcc4a35aa45bc2858756533683a159ac1f44ec4fed5
+size 2407338

resources/infoDict.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5602ffd87f6bb46b53ac6c4cece0703a3595721a1b8cc6cc4e650154bdbfd6e
+size 10577505

resources/labelSet.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f46f078cdab57f85ac70a3da7e11cc9f001c493a40e1ea32c7da50acfb8af2b
+size 36457

resources/leafImageIndices.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26463c81aa1fdde63490deafb70016d6efe5d5241d25f4e4dd63aeac14b4452b
+size 144737

resources/leafImageLabels.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2f8ae09eb663be008a260a2f30d668ff28b042a72578178783017add5a5034f7
+size 259577

resources/leafLabelSet.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c25e4e98ea2bd5a313ca032c212aeb688e84afca6458757889be84a282c4580e
+size 25482

resources/leafMeansAndStds.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ab072ba014fe39029378b2c088d8a8ab85f1291163e25e060e621f00cedb194c
+size 236

resources/leafspeciesIndexDict.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92fbc241dc2825a4f83beb61e3be3a4ae59e0d04cb3a68a80e58fbb3aff0d4ec
+size 69399

resources/meansAndStds.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f9c146afe67291a4143714b466a3cb7eed425bfed1218103930ebafeefc0747
+size 236

resources/speciesNameToID.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ffb83cf9e6053dbc3888d5e6b3e2d784bfd09333fb5058b51f3f1da0d005fc1
+size 584518

resources/speciesNameToKey.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0edd05755d1b9ca01291426ff35a9964589d4cd52f1cbea3f032b9dbfc7f2b17
+size 94428

resources/speciesNameToVernacular.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2207518c0f274b0b579fe8bf70c6605b3abaeb7a725adebab5b2bd3ec6201e41
+size 140185